/***************************** 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 <iomanip>

#include "MvException.h"
#include "MvMiscellaneous.h"
#include "MvSci.h"
#include "MvNetCDF.h"
#include "Thermo.h"

const double ZEPQMI = 1.0e-12;  // minimum threshold for specific humidity
const int XWIND_POS = 1010;     // x wind position

//_____________________________________________________________________________

MvThermoDataBase::MvThermoDataBase(MvField& fieldToken, const ThermoParams& pars) :
    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
    std::vector<long> values_ndim;
    std::vector<std::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
    std::vector<double> xv(size, XWIND_POS);
    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_, false);  //-- 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(std::max(ZEPQMI, val));
        dCount_++;
    }
    else if (par == pars_->lnPressId_)
        processPressure(val);
    else {
        std::ostringstream oss;
        oss << "Thermo-> Field contains unknown data, param = " << par;
        throw MvException(oss.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)
    std::ostringstream oss;
    oss << std::setfill('0')
        << std::setw(8) << date
        << std::setw(4) << (time ? time : "TTTT")
        << std::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_;
    std::vector<long> values_ndim;
    std::vector<std::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: {
            std::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_));
            std::ostringstream oss;
            oss << loc.latitude() << '/'
                << loc.longitude();

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

        case eAreaAverage: {
            std::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;
}

std::string
MvThermoDataBase::formulaPrintName()
{
    std::string name = pars_->dewPointFormulaType_;
    for (char& it : name) {
        char cc = it;
        it = (cc == '_') ? ' ' : tolower(cc);
    }

    return name;
}

std::string
MvThermoDataBase::title(MvNetCDF& netcdf)
{
    // Get first time information
    std::string sdate = timeKey_.substr(0, 8);
    std::string stime = timeKey_.substr(8, 4);
    int istep = atoi(timeKey_.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");
            std::string 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 std::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(MvSci::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(MvSci::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(MvSci::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(std::max(ZEPQMI, val));
        dCount_++;
    }
    else if (par == pars_->lnPressId_)
        processPressure(val);
    else {
        std::ostringstream oss;
        oss << "Thermo-> field contains unknown data, param = " << par;
        throw MvException(oss.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_) {
        std::ostringstream oss;
        oss << "Thermo-> UKmodel, the numbers of fields are inconsistent = " << pCount_ << " " << tCount_ << " " << dCount_;
        throw MvException(oss.str());
    }

    // Compute and save levels and dewpoints values
    for (int k = 0; k < dCount_; k++) {
        TD_.push_back(MvSci::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())
        return new MvThermoDataHalfLevels(f, p);

    //-- non-ECMWF data
    if (f.vertCoordCoefPairCount() == 1)
        return new MvThermoDataFullLevels(f, p);

    //-- pressure level data
    if (f.isPressureLevel())
        return new MvThermoDataPressureLevels(f, p);

    //-- UKMO New Dynamics data
    if (f.levelType() == cML_UKMO_ND)
        return new MvThermoDataUkmoNDLevels(f, p);

    return nullptr;  //-- unimplemented data
}
