/***************************** 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 <Metview.h>
#include <MvException.h>
#include "MvSci.h"

#include <cassert>
#include <map>

class ThetaBase : public MvService
{
public:
    enum LevelType
    {
        NoLevelType,
        PressureLevelType,
        ModelLevelType
    };
    enum ParamIds
    {
        TemperatureId = 130,
        SpecHumId = 133,
        LnspId = 152,
        RelHumId = 157
    };

    ThetaBase(const char* a, LevelType levelType, int resultParamId) :
        MvService(a), levelType_(levelType), resultParamId_(resultParamId)
    {
        if (levelType_ == PressureLevelType) {
            strcpy(Plevtype_, "PL");
        }
        else if (levelType_ == ModelLevelType) {
            strcpy(Plevtype_, "ML");
        }
    }
    void serve(MvRequest&, MvRequest&) override;

protected:
    virtual MvRequest process(MvRequest& In) = 0;
    void checkGrid(MvField& f, int fieldIdx);
    void checkValNum(MvField& f, int fieldIdx, int refValNum);
    void checkPressure(MvField& f, int fieldIdx, double p);
    std::string formatFieldStr(MvField& f, int fieldIdx);

    char Plevtype_[10];
    LevelType levelType_;
    int resultParamId_;
    MvFilter filterPar_{"PARAM"};
    MvFilter filterLevType_{"LEVTYPE"};
};

class ThetaNoHum : public ThetaBase
{
public:
    ThetaNoHum(const char* a, ThetaBase::LevelType levelType, int resultParamId) :
        ThetaBase(a, levelType, resultParamId) {}

protected:
    MvRequest process(MvRequest&) override;
    MvRequest compute(MvFieldSetIterator& tIter, MvField& lnspF);
    virtual void computeVals(MvField& tF, MvField& lnsp, double p, int level, bool flag) = 0;
};

class ThetaHum : public ThetaBase
{
public:
    ThetaHum(const char* a, ThetaBase::LevelType levelType, int resultParamId) :
        ThetaBase(a, levelType, resultParamId) {}

protected:
    MvRequest process(MvRequest&) override;
    MvRequest compute(int num, MvFieldSetIterator& tIter, MvFieldSetIterator& qIter, MvField& lnspF);
    virtual void computeVals(MvField& tF, MvField& qF, MvField& lnsp, double p, int level, bool flag) = 0;
};

class ThetaPl : public ThetaNoHum
{
public:
    ThetaPl(const char* a) :
        ThetaNoHum(a, ThetaBase::PressureLevelType, 3) {}

protected:
    void computeVals(MvField& tF, MvField& lnsp, double p, int level, bool flag) override;
};

class ThetaMl : public ThetaNoHum
{
public:
    ThetaMl(const char* a) :
        ThetaNoHum(a, ThetaBase::ModelLevelType, 3) {}

protected:
    void computeVals(MvField& tF, MvField& lnsp, double p, int level, bool flag) override;
};

class ThetaEqPl : public ThetaHum
{
public:
    ThetaEqPl(const char* a) :
        ThetaHum(a, ThetaBase::PressureLevelType, 4) {}

protected:
    void computeVals(MvField& tF, MvField& qF, MvField& lnsp, double p, int level, bool flag) override;
};

class ThetaEqMl : public ThetaHum
{
public:
    ThetaEqMl(const char* a) :
        ThetaHum(a, ThetaBase::ModelLevelType, 4) {}

protected:
    void computeVals(MvField& tF, MvField& qF, MvField& lnsp, double p, int level, bool flag) override;
};

class ThetaEqSatPl : public ThetaNoHum
{
public:
    ThetaEqSatPl(const char* a) :
        ThetaNoHum(a, ThetaBase::PressureLevelType, 5) {}

protected:
    void computeVals(MvField& tF, MvField& lnsp, double p, int level, bool flag) override;
};

class ThetaEqSatMl : public ThetaNoHum
{
public:
    ThetaEqSatMl(const char* a) :
        ThetaNoHum(a, ThetaBase::ModelLevelType, 5) {}

protected:
    void computeVals(MvField& tF, MvField& lnsp, double p, int level, bool flag) override;
};

//===========================================================
//
//  TheataBase
//
//===========================================================

void ThetaBase::serve(MvRequest& in, MvRequest& out)
{
    try {
        out = process(in);
    }
    catch (MvException& e) {
        setError(1, "POTTF-> %s", e.what());
    }
}

void ThetaBase::checkGrid(MvField& f, int fieldIdx)
{
    if (f.isSpectral())
        throw MvException(
            "Spectral representation is not supported, " +
            formatFieldStr(f, fieldIdx) + "!");
}

void ThetaBase::checkValNum(MvField& f, int fieldIdx, int refValNum)
{
    if (f.countValues() != refValNum)
        throw MvException("Different number of points=" +
                          std::to_string(f.countValues()) +
                          " in field than expected=" +
                          std::to_string(refValNum) + " " +
                          formatFieldStr(f, fieldIdx) + "!");
}

void ThetaBase::checkPressure(MvField& f, int fieldIdx, double p)
{
    if (p <= 0 || p > 1500. * 100.)
        throw MvException(
            "Invalid pressure value=" + std::to_string(p) +
            ", " + formatFieldStr(f, fieldIdx));
}

std::string ThetaBase::formatFieldStr(MvField& f, int fieldIdx)
{
    return "field=[param=" + f.parameterName() + ",index=" + std::to_string(fieldIdx) + "]";
}

//===========================================================
//
//  TheataNoHum
//
//===========================================================

MvRequest ThetaNoHum::process(MvRequest& in)
{
    MvRequest out;
    MvRequest tGrib;
    in.getValue(tGrib, "TEMPERATURE");
    MvFieldSet tFs(tGrib);
    MvFieldSetIterator tIter(tFs);
    tIter.setFilter(filterPar_ == TemperatureId && filterLevType_ == Plevtype_);

    if (levelType_ == PressureLevelType) {
        MvField lnsp;
        out = compute(tIter, lnsp);
    }
    else if (levelType_ == ModelLevelType) {
        MvRequest lnspGrib;
        in.getValue(lnspGrib, "LNSP");
        MvFieldSet lnspFs(lnspGrib);
        MvFieldSetIterator lnspIter(lnspFs);
        lnspIter.setFilter(filterPar_ == LnspId && filterLevType_ == Plevtype_);
        {
            MvField lnspF = lnspIter();
            if (!lnspF) {
                throw("No lnsp field is found in input!");
            }
            MvFieldExpander xp_lnsp(lnspF);
            out = compute(tIter, lnspF);
        }
    }

    return out;
}

MvRequest ThetaNoHum::compute(MvFieldSetIterator& tIter, MvField& lnspF)
{
    int valNum = 0;
    if (levelType_ == ModelLevelType) {
        valNum = lnspF.countValues();
    }

    MvFieldSet out;
    MvField tF;

    int fieldCnt = 0;
    while ((tF = tIter())) {
        MvFieldExpandThenFree xpt(tF);
        checkGrid(tF, fieldCnt);
        double p = 0.;
        int level = 1;
        if (levelType_ == PressureLevelType) {
            p = tF.level() * 100;
            checkPressure(tF, fieldCnt, p);
        }
        else if (levelType_ == ModelLevelType) {
            checkValNum(tF, fieldCnt, valNum);
            level = static_cast<int>(tF.level());
        }

        bool humCompFlag = (tF.yyyymmddFoh() < 19830421.);
        computeVals(tF, lnspF, p, level, humCompFlag);

        tF.setGribKeyValueLong("paramId", resultParamId_);
        out += tF;
        out[fieldCnt].setShape(packed_mem);
        fieldCnt++;
    }

    return out.getRequest();
}

//===========================================================
//
//  TheataHum
//
//===========================================================

MvRequest ThetaHum::process(MvRequest& in)
{
    MvRequest out;
    MvRequest tGrib;
    in.getValue(tGrib, "TEMPERATURE");
    MvFieldSet tFs(tGrib);
    MvFieldSetIterator tIter(tFs);
    tIter.setFilter(filterPar_ == TemperatureId && filterLevType_ == Plevtype_);

    MvRequest qGrib;
    in.getValue(qGrib, "HUMIDITY");
    MvFieldSet qFs(qGrib);
    MvFieldSetIterator qIter(qFs);
    qIter.setFilter((filterPar_ == SpecHumId || filterPar_ == RelHumId) && filterLevType_ == Plevtype_);

    int tNum = 0;
    int qNum = 0;
    while (tIter()) {
        tNum++;
    }
    tIter.rewind();
    while (qIter()) {
        qNum++;
    }
    qIter.rewind();

    if (tNum == 0)
        throw MvException("No temperature field is found in input");

    if (qNum == 0)
        throw MvException("No humidity field is found in input");

    if (tNum != qNum)
        throw MvException("Different number of temperature (=" + std::to_string(tNum) +
                          ") and humidity (=" + std::to_string(qNum) + ") fields!");

    if (levelType_ == PressureLevelType) {
        MvField lnspF;
        out = compute(tNum, tIter, qIter, lnspF);
    }
    else if (levelType_ == ModelLevelType) {
        MvRequest lnspGrib;
        in.getValue(lnspGrib, "LNSP");
        MvFieldSet lnspFs(lnspGrib);
        MvFieldSetIterator lnspIter(lnspFs);
        lnspIter.setFilter(filterPar_ == LnspId && filterLevType_ == Plevtype_);
        {
            MvField lnspF = lnspIter();
            if (!lnspF) {
                throw("No lnsp field is found in input!");
            }
            MvFieldExpander xp_lnsp(lnspF);
            out = compute(tNum, tIter, qIter, lnspF);
        }
    }

    return out;
}

MvRequest ThetaHum::compute(int fNum, MvFieldSetIterator& tIter, MvFieldSetIterator& qIter, MvField& lnspF)
{
    int valNum = 0;
    if (levelType_ == ModelLevelType) {
        valNum = lnspF.countValues();
    }

    MvFieldSet out;
    double eps = 1E-6;

    for (int i = 0; i < fNum; i++) {
        MvField tF = tIter();
        MvField qF = qIter();
        MvFieldExpandThenFree xpt(tF);
        MvFieldExpandThenFree xpq(qF);

        checkGrid(tF, i);
        checkGrid(qF, i);

        double p = 0.;
        int level = 1;
        if (levelType_ == PressureLevelType) {
            p = tF.level() * 100.;
            checkPressure(tF, i, p);
            double pFromQ = qF.level() * 100.;
            if (fabs(p - pFromQ) > eps) {
                throw MvException("Pressure level does not match for temperature and humidity, field [" + std::to_string(i) + "]!");
            }
            valNum = tF.countValues();
            double valNumFromQ = qF.countValues();
            if (valNum != valNumFromQ) {
                throw MvException("Temperature and humidity does not have the same numbe rof gridpoints, field [" + std::to_string(i) + "]!");
            }
        }
        else if (levelType_ == ModelLevelType) {
            checkValNum(tF, i, valNum);
            checkValNum(qF, i, valNum);
            level = static_cast<int>(tF.level());
        }

        bool humCompFlag = (tF.yyyymmddFoh() < 19830421.);
        computeVals(tF, qF, lnspF, p, level, humCompFlag);

        tF.setGribKeyValueLong("paramId", resultParamId_);
        out += tF;
        out[i].setShape(packed_mem);
    }

    return out.getRequest();
}

//===========================================================
//
//  TheataPl
//
//===========================================================

void ThetaPl::computeVals(MvField& tF, MvField& /*lnsp*/, double p, int /*level*/, bool /*flag*/)
{
    double pTerm = MvSci::potentialTemperaturePressureTerm(p);
    for (int i = 0; i < tF.countValues(); i++) {
        if (!MISSING_VALUE(tF[i])) {
            tF[i] = tF[i] * pTerm;
        }
        else {
            tF.setValueAtToMissing(i);
        }
    }
}

//===========================================================
//
//  TheataMl
//
//===========================================================

void ThetaMl::computeVals(MvField& tF, MvField& lnspF, double p, int level, bool /*flag*/)
{
    for (int i = 0; i < tF.countValues(); i++) {
        p = lnspF.meanML_to_Pressure_byLNSP(lnspF[i], level);
        if (!MISSING_VALUE(tF[i]) && !MISSING_VALUE(p)) {
            tF[i] = MvSci::potentialTemperature(tF[i], p);
        }
        else {
            tF.setValueAtToMissing(i);
        }
    }
}

//===========================================================
//
//  TheataEqPl
//
//===========================================================

void ThetaEqPl::computeVals(MvField& tF, MvField& qF, MvField& /*lnspF*/, double p, int /*level*/, bool flag)
{
    bool hasRelHum = (static_cast<int>(qF.parameter()) == RelHumId);

    for (int i = 0; i < tF.countValues(); i++) {
        if (!MISSING_VALUE(tF[i]) && !MISSING_VALUE(qF[i])) {
            double q = qF[i];
            if (hasRelHum) {
                q = q * MvSci::saturationSpecHumidity(tF[i], p, flag) * 0.01;
            }
            tF[i] = MvSci::equivalentPotentialTemperature(tF[i], p, q, flag);
        }
        else {
            tF.setValueAtToMissing(i);
        }
    }
}

//===========================================================
//
//  TheataEqMl
//
//===========================================================

void ThetaEqMl::computeVals(MvField& tF, MvField& qF, MvField& lnspF, double p, int level, bool flag)
{
    bool hasRelHum = (static_cast<int>(qF.parameter()) == RelHumId);

    for (int i = 0; i < tF.countValues(); i++) {
        p = lnspF.meanML_to_Pressure_byLNSP(lnspF[i], level);
        if (!MISSING_VALUE(tF[i]) && !MISSING_VALUE(qF[i]) && !MISSING_VALUE(p)) {
            double q = qF[i];
            if (hasRelHum) {
                q = q * MvSci::saturationSpecHumidity(tF[i], p, flag) * 0.01;
            }
            tF[i] = MvSci::equivalentPotentialTemperature(tF[i], p, q, flag);
        }
        else {
            tF.setValueAtToMissing(i);
        }
    }
}

//===========================================================
//
//  TheataEqSatPl
//
//===========================================================

void ThetaEqSatPl::computeVals(MvField& tF, MvField& /*lnsp*/, double p, int /*level*/, bool flag)
{
    double pTerm = MvSci::potentialTemperaturePressureTerm(p);
    for (int i = 0; i < tF.countValues(); i++) {
        if (!MISSING_VALUE(tF[i])) {
            tF[i] = MvSci::saturatedEquivalentPotentialTemperatureWithPTerm(tF[i], p, pTerm, flag);
        }
        else {
            tF.setValueAtToMissing(i);
        }
    }
}

//===========================================================
//
//  TheataEqSatMl
//
//===========================================================

void ThetaEqSatMl::computeVals(MvField& tF, MvField& lnspF, double p, int level, bool flag)
{
    for (int i = 0; i < tF.countValues(); i++) {
        p = lnspF.meanML_to_Pressure_byLNSP(lnspF[i], level);
        if (!MISSING_VALUE(tF[i]) && !MISSING_VALUE(p)) {
            tF[i] = MvSci::saturatedEquivalentPotentialTemperature(tF[i], p, flag);
        }
        else {
            tF.setValueAtToMissing(i);
        }
    }
}


int main(int argc, char** argv)
{
    MvApplication theApp(argc, argv);
    ThetaMl thetaMl("POTT_M");
    ThetaPl thetaPl("POTT_P");
    ThetaEqMl eqpottMl("EQPOTT_M");
    ThetaEqPl eqpottPl("EQPOTT_P");
    ThetaEqSatMl seqpottMl("SEQPOTT_M");
    ThetaEqSatPl seqpottPL("SEQPOTT_P");

    theApp.run();
}
