/***************************** LICENSE START ***********************************

 Copyright 2012 ECMWF and INPE. This software is distributed under the terms
 of the Apache License version 2.0. In applying this license, ECMWF does not
 waive the privileges and immunities granted to it by virtue of its status as
 an Intergovernmental Organization or submit itself to any jurisdiction.

 ***************************** LICENSE END *************************************/

#include "DataThermo.h"
#include "MvException.h"
#include "MvMiscelaneous.h"
#include "MvNetCDF.h"
#include "Thermo.h"

const int    cFREE  = -999;    // to mark free array cells
const double ZEPQMI = 1.0e-12; // minimum threshold for specific humidity

//_____________________________________________________________________________

MvThermoDataBase::MvThermoDataBase( MvField& fieldToken, const ThermoParams& pars ) :
   lnspFlag_(false),
   counter_( 0 ),
   tCount_( 0 ),
   dCount_( 0 ),
   uCount_( 0 ),
   vCount_( 0 ),
   pCount_( 0 ),
   nTIME_STR_( 17 ),
   currField_( fieldToken ),
   pars_(&pars)
{
}

void
MvThermoDataBase::initialiseDataMemory()
{
   uCount_ = vCount_ = tCount_ = dCount_ = pCount_ = 0;

   T_.clear(); Q_.clear();  TD_.clear();
   P_.clear(); PUV_.clear();
   U_.clear(); V_.clear();
}

void
MvThermoDataBase::readOneSetOfFields( MvFieldSetIterator& iter )
{
   const double cBIGDOUBLE = 3.4E35;
   const double oneMinute = 1.0 / 24.0;

   // Get time key info
   this->generateTimeKey();

   // Initialise data values storage
   this->initialiseDataMemory();

   // Main loop
   // It assumes that the input fieldset is sorted by Levelist, e.g.,
   // fieldset = Level1 Param1..n, Level2 Param1...n, ..., Levelm Param1...n
   bool flag  = true;
   while( flag )    //-- read while same valid time
   {
     // Get value from the field
      double val = this->getValueFromCurrentField();
      if ( val > cBIGDOUBLE )
         throw MvException("Thermo: Unsupported grid or point outside field!");

      // Save data and pressure values to the appropriated structure
      this->collectValue(val);

      // Move to the next field
      nextField_ = iter();
      if( ! nextField_ )
         flag = false;                //-- end-of-file
      else
         flag = ( fabs(currField_.yyyymmddFoh() - nextField_.yyyymmddFoh()) < oneMinute ) &&
                ( fabs(currField_.stepFoh()     - nextField_.stepFoh())     < oneMinute );

      if( flag )
         currField_ = nextField_;
    }

   // Set of current fields preprocessed - everything ok so far?
   if( tCount_ != dCount_ )
      throw MvException("Thermo: Number of levels for T, Q are inconsistent!");

   if( tCount_ <= 1 )
       throw MvException("Thermo: T fields are missing!");

   if( needsLnsp() && !lnspFlag_ )
       throw MvException("Thermo: LNSP field missing!");

   if( uCount_ != vCount_ )
      throw MvException("Thermo: Different number of levels for U and V!");
}

bool
MvThermoDataBase::generate( MvNetCDF& netcdf )
{
   // Create netcdf variables, dimensions and global atributes
   if ( counter_ == 0 )
   {
      if ( !this->ncInit(netcdf) )
         throw MvException("Thermo: Failure to initilise a netCDF file!");
   }

   // Get netcdf variables
   // tCount_ and dCount_ should have the same value
   long size = tCount_;
   MvNcVar *time = netcdf.getVariable(sTIME);
   MvNcVar *temp = netcdf.getVariable(sTEMP);
   MvNcVar *dew = netcdf.getVariable(sDEWPOINT);
   MvNcVar *press = netcdf.getVariable(sPRESSURE);
   if ( !time || !temp || !dew || !press )
      throw MvException("Thermo: ERROR accessing netCDF variables");

   // Compute pressure and dewpoint
   this->computeLevelsTd();

   // Re-scale values
   this->rescaleValues();

   // Add data to the variables
   const char* key = timeKey_.c_str();
   time->setCurrent(counter_);
   if (!time->put(key,1,strlen(key)+1))
      throw MvException("Thermo: ERROR writing a netCDF variable: time");

   temp->setCurrent(counter_);
   if (!temp->put(T_,1,size))
      throw MvException("Thermo: ERROR writing a netCDF variable: t");

   dew->setCurrent(counter_);
   if (!dew->put(TD_,1,size))
      throw MvException("Thermo: ERROR writing a netCDF variable: td");

   press->setCurrent(counter_);
   if (!press->put(P_,1,size))
      throw MvException("Thermo: ERROR writing a netCDF variable: pres");

   // Add wind info to netcdf
   if( windExist() )
      if ( !this->generateWind(netcdf) )
         return false;

   counter_++;

   return true;
}

bool
MvThermoDataBase::generateWind(MvNetCDF &netcdf)
{
   // Number of levels
   long size = uCount_;  // uCount_ and vCount_ should have the same value

   // Generate NetCDF variables
   vector<long> values_ndim;
   vector<string> values_sdim;
   values_ndim.push_back(0);
   values_ndim.push_back(size);
   values_sdim.push_back(sTIME);
   values_sdim.push_back(sNLEVWIND);

   MvNcVar *uwind = netcdf.addVariable(sU,ncDouble,values_ndim,values_sdim);
   MvNcVar *vwind = netcdf.addVariable(sV,ncDouble,values_ndim,values_sdim);
   MvNcVar *xvals = netcdf.addVariable(sXWIND,ncDouble,values_ndim,values_sdim);
   MvNcVar *pwind = netcdf.addVariable(sPWIND,ncDouble,values_ndim,values_sdim);

   // Add U values
   uwind->setCurrent(counter_);
   if (!uwind->put(U_,1,size))
      throw MvException("Thermo: ERROR writing a netCDF variable: u");

   // Add V values
   vwind->setCurrent(counter_);
   if (!vwind->put(V_,1,size))
      throw MvException("Thermo: ERROR writing a netCDF variable: v");

   // Add horizontal axis values
   vector <double> xv(size,1000.5);
   xvals->setCurrent(counter_);
   if (!xvals->put(xv,1,size))
      throw MvException("Thermo: ERROR writing a netCDF variable: xwind");

   // Add pressure wind values
   pwind->setCurrent(counter_);
   if (!pwind->put(PUV_,1,size))
      throw MvException("Thermo: ERROR writing a netCDF variable: pwind");

   return true;
}

//
//  Value can be extracted in several ways
//
double
MvThermoDataBase::getValueFromCurrentField()
{
   double y;
   switch( pars_->pointType_ )
   {
      case eInterpolated:
         y = currField_.interpolateAt( pars_->lon_, pars_->lat_ );    //-- x/y
         break;

      case eGridPoint:
         y = currField_.nearestGridpoint( pars_->lon_, pars_->lat_ ); //-- x/y
         break;

      case eAreaAverage:
         y = currField_.integrate( pars_->lat_, pars_->lon_, pars_->lat2_, pars_->lon2_ ); //-- n/w/s/e
         break;
   }

   return y;
}

void
MvThermoDataBase::collectValue( double val )
{
   // Get param and level
   int par = currField_.parameter();
   int level = this->getLevel();

   // Save param and level values
   // Levels for wind and temp/dewpoint are only saved once
   if( par == pars_->uFieldId_ )
   {
      U_.push_back(val);
      uCount_++;
   }
   else if( par == pars_->vFieldId_ )
   {
      V_.push_back(val);
      PUV_.push_back(level);
      vCount_++;
   }
   else if( par == pars_->tempId_ )
   {
      T_.push_back(val);
      P_.push_back(level);
      tCount_++;
   }
   else if( par == pars_->dewId_ )
   {
      Q_.push_back(max(ZEPQMI,val));
      dCount_++;
   }
   else if( par == pars_->lnPressId_ )
      processPressure( val );
   else
   {
      ostringstream oss;
      oss << "Tephigrams: field contains unknown data, param = " << par;
      throw MvException( oss.str().c_str() );
   }
}

void
MvThermoDataBase::rescaleValues()
{
   for ( int i = 0; i < tCount_; i++ )
   {
      T_[i]  = T_[i]  - K2C;
      TD_[i] = TD_[i] - K2C;
      P_[i]  = P_[i]  / 100.;
   }

   // Re-scale pressure related to the wind
   if ( windExist() )
   {
      for ( int i = 0; i < uCount_; i++ )
         PUV_[i]  = PUV_[i]  / 100.;
   }
}

void
MvThermoDataBase::generateTimeKey()
{
   MvRequest    req    = currField_.getRequest();

   const char*  date   = req("DATE");
   const char*  time   = req("TIME");
   const char*  step   = req("STEP");

   // variable nTIME_STR = 17 (8+4+4+1)
   ostringstream oss;
   oss << setfill( '0' )
       << setw( 8 ) << date
       << setw( 4 ) << ( time ? time : "TTTT" )
       << setw( 4 ) << ( step ? step : "SSSS" );

   timeKey_ = oss.str();
}

bool MvThermoDataBase::ncInit(MvNetCDF &netcdf)
{
   // Add global attributes to the netcdf file
   if ( !this->ncGlobalAttributes(netcdf) )
      return false;

   // Create variables
   long size = tCount_;
   vector<long> values_ndim;
   vector<string> values_sdim;
   values_ndim.push_back(0);
   values_sdim.push_back(sTIME);

   values_ndim.push_back(nTIME_STR_);
   values_sdim.push_back(sTIME_STR);
   if ( !netcdf.addVariable(sTIME,ncChar,values_ndim,values_sdim) )
      return false;

   values_ndim[1] = size;
   values_sdim[1] = sNLEV;
   if ( !netcdf.addVariable(sTEMP,ncFloat,values_ndim,values_sdim) )
      return false;
   if ( !netcdf.addVariable(sDEWPOINT,ncFloat,values_ndim,values_sdim) )
      return false;
   if ( !netcdf.addVariable(sPRESSURE,ncFloat,values_ndim,values_sdim) )
      return false;

   return true;
}

bool MvThermoDataBase::ncGlobalAttributes(MvNetCDF &netcdf)
{
   netcdf.addAttribute("_FILL_VALUE",MissingValue);
   netcdf.addAttribute("DewPoint Formulation",formulaPrintName().c_str());
   switch( pars_->pointType_ )
   {
      case eInterpolated:
      {
         ostringstream oss;
         oss << pars_->lat_
             << '/'
             << pars_->lon_;

         return netcdf.addAttribute("Coordinates",oss.str().c_str());
      }
      break;

      case eGridPoint:
      {
         MvLocation loc = currField_.nearestGridPointLocation( MvLocation( pars_->lat_, pars_->lon_ ) );
         ostringstream oss;
         oss << loc.latitude() << '/'
             << loc.longitude();

         return netcdf.addAttribute("Coordinates",oss.str().c_str());
      }
      break;

      case eAreaAverage:
      {
         ostringstream oss;
         oss << pars_->lat_  << '/'
             << pars_->lon_  << '/'
             << pars_->lat2_ << '/'
             << pars_->lon2_;

         return netcdf.addAttribute("Coordinates",oss.str().c_str());
      }
      break;
   }

   return true;
}

bool
MvThermoDataBase::windExist()
{
   bool flag = uCount_ ? true : false;
   return flag;
}

string
MvThermoDataBase::formulaPrintName()
{
   string name = pars_->dewPointFormulaType_;
   for( unsigned int it = 0; it < name.length(); ++it )
   {
      char cc = name[it];
      name[it] = ( cc == '_' ) ? ' ' : tolower(cc);
   }

   return name;
}

string
MvThermoDataBase::title( MvNetCDF& netcdf)
{
   // Get first time information from the netcdf file
   MvNcVar *time = netcdf.getVariable(sTIME);
   NcDim* dim = time->getDimension(1);
   int size = dim->size();
   vector <char> vtime;
   time->setCurrent(0L);
   time->get(vtime,1,size);

   // Decode time info
   string str(vtime.begin(),vtime.end());
   string sdate = str.substr(0,8);
   string stime = str.substr(8,4);
   int istep    = atoi(str.substr(12,4).c_str());

   // Create title
   char  title[200];
   switch( pars_->pointType_ )
   {
      case eInterpolated:
         sprintf(title,"%s %s step %d [%.2f,%.2f] %s",
                 sdate.c_str(), stime.c_str(), istep,
                 pars_->lat_, pars_->lon_,
                 this->formulaPrintName().c_str() );
         break;

      case eGridPoint:
      {
         // Coordinates used may not be the same as the input ones
         MvNcAtt* tmpatt = netcdf.getAttribute("Coordinates");
         str = tmpatt->as_string(0);
         int ipos = str.find('/');
         float lat = atof(str.substr(0,ipos).c_str());
         float lon = atof(str.substr(ipos+1).c_str());
         sprintf(title,"%s %s step %d [nearest: %.2f,%.2f] %s",
                 sdate.c_str(), stime.c_str(), istep, lat, lon,
                 this->formulaPrintName().c_str() );
      }
      break;

      case eAreaAverage:
         sprintf(title,"%s %s step %d [%.1f,%.1f,%.1f,%.1f] %s",
                 sdate.c_str(), stime.c_str(), istep,
                 pars_->lat_, pars_->lon_, pars_->lat2_, pars_->lon2_,
                 this->formulaPrintName().c_str() );
      break;
   }

   return string(title);
}

//==============================================================
// MvThermoDataHalfLevels
//______________________________________________________________

void
MvThermoDataHalfLevels::computeLevelsTd()
{
   for( int k = 0; k < dCount_; k++)
   {
      P_[k] = currField_.meanML_to_Pressure_bySP( exp(LNSP_), int(P_[k]) );
      TD_.push_back( metview::dewPointFromQ( P_[k], T_[k], Q_[k], pars_->dewPointFormulaType_ ) );
   }

   // Wind: model to pressure levelIndex
   if ( windExist() )
   {
      for( int k = 0; k < uCount_; k++)
         PUV_[k] = currField_.meanML_to_Pressure_bySP( exp(LNSP_), int(PUV_[k]) );
   }
}

void
MvThermoDataHalfLevels::processPressure( double val )
{
   lnspFlag_ = true;;
   LNSP_ = val;
}


//=============================================================
// MvThermoDataFullLevels
//_____________________________________________________________

void
MvThermoDataFullLevels::computeLevelsTd()
{
   for( int k = 0; k < dCount_; k++)
   {
      P_[k] = currField_.ML_to_Pressure_bySP( exp(LNSP_), int(P_[k]) );
      TD_.push_back( metview::dewPointFromQ( P_[k], T_[k], Q_[k], pars_->dewPointFormulaType_ ) );
   }
}

//_____________________________________________________________________________

void
MvThermoDataFullLevels::processPressure( double val )
{
   lnspFlag_ = true;
   LNSP_ = val;
}

//==============================================================
// MvThermoDataPressureLevels
//______________________________________________________________

MvThermoDataPressureLevels::MvThermoDataPressureLevels( MvField& fieldToken, const ThermoParams& pars )
    : MvThermoDataBase( fieldToken, pars )
{
}

void
MvThermoDataPressureLevels::computeLevelsTd()
{
   for( int k = 0; k < dCount_; k++)
   {
      TD_.push_back( metview::dewPointFromQ( P_[k]*100., T_[k], Q_[k], pars_->dewPointFormulaType_) );
   }
}

void
MvThermoDataPressureLevels::processPressure( double )
{
  //-- not needed if pressure levels
}

void
MvThermoDataPressureLevels::rescaleValues()
{
   for ( int i = 0; i < tCount_; i++ )
   {
      T_[i]  = T_[i]  - K2C;
      TD_[i] = TD_[i] - K2C;
   }
}

//==============================================================
// MvThermoDataUkmoNDLevels
//______________________________________________________________

MvThermoDataUkmoNDLevels::MvThermoDataUkmoNDLevels( MvField& fieldToken, const ThermoParams& pars )
    : MvThermoDataBase( fieldToken, pars )
{
}

void
MvThermoDataUkmoNDLevels::collectValue( double val )
{
   // Get param and level
   int par = currField_.parameter();
   int level = this->getLevel();

   // Save param and level values
   // Levels for wind are only saved once
   if( par == pars_->uFieldId_ )
   {
      U_.push_back(val);
      uCount_++;
   }
   else if( par == pars_->vFieldId_ )
   {
      V_.push_back(val);
      PUV_.push_back(level);
      vCount_++;
   }
   else if( par == pars_->tempId_ )
   {
      T_.push_back(val);
      tCount_++;
   }
   else if( par == pars_->dewId_ )
   {
      Q_.push_back(max(ZEPQMI,val));
      dCount_++;
   }
   else if( par == pars_->lnPressId_ )
      processPressure( val );
   else
   {
      ostringstream oss;
      oss << "Tephigrams: field contains unknown data, param = " << par;
      throw MvException( oss.str().c_str() );
   }
}

void
MvThermoDataUkmoNDLevels::processPressure( double val)
{
   pCount_++;
   P_.push_back(val);
}

void
MvThermoDataUkmoNDLevels::computeLevelsTd()
{
   // Consistency check 1: assumes that all input data are related to the
   // same set of levelist. Therefore, all the data counters should have
   // the same value.
   if ( pCount_ != tCount_ || pCount_ != dCount_ )
   {
      ostringstream oss;
      oss << "Thermo: UKmodel, the numbers of fields are inconsistent = " << pCount_ << " " << tCount_ << " " << dCount_;
      throw MvException( oss.str().c_str() );
   }

   // Compute and save levels and dewpoints values
   for( int k = 0; k < dCount_; k++)
   {
      TD_.push_back(metview::dewPointFromQ( P_[k], T_[k], Q_[k], pars_->dewPointFormulaType_ ));
   }
}

void
MvThermoDataUkmoNDLevels::rescaleValues()
{
   for ( int i = 0; i < tCount_; i++ )
   {
      T_[i]  = T_[i]  - K2C;
      TD_[i] = TD_[i] - K2C;
   }
}

//________________________________________________________
//
//  Factory to create and return an object of the required
//  tephidata class.
//
MvThermoDataBase* MvThermoDataFactory( MvField& f, const ThermoParams& p )
{
   const int cECMWF = 98;

   MvFieldExpander expa( f );

   //-- "standard" ECMWF data
   if( f.centre() == cECMWF && f.isModelLevel() )
   {
      cout << "Factory: MvThermoDataHalfLevels" << endl;
      return new MvThermoDataHalfLevels( f, p );
   }

   //-- non-ECMWF data
   if( f.vertCoordCoefPairCount() == 1 )
   {
      cout << "Factory: MvThermoDataFullLevels" << endl;
      return new MvThermoDataFullLevels( f, p );
   }

   //-- pressure level data
   if( f.isPressureLevel() )
   {
      cout << "Factory: MvThermoDataPressureLevels" << endl;
      return new MvThermoDataPressureLevels( f, p );
   }

   //-- UKMO New Dynamics data
   if( f.levelType() == cML_UKMO_ND )
   {
      cout << "Factory: MvThermoDataUkmoNDLevels" << endl;
      return new MvThermoDataUkmoNDLevels( f, p );
   }

   return 0;  //-- unimplemented data
}
