
/***************************** 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 <string>
#include <vector>
#include <iostream>
#include <cstdio>
#include <cstring>
#include <cassert>

#include "grib_api.h"

#include "MvFlexpart.h"
#include "MvBinaryReader.h"
#include "MvDate.h"
#include "MvMiscellaneous.h"

//#define _GNU_SOURCE
//#include <fenv.h>
// int feenableexcept(int excepts);


class MvFlexpartGrid
{
public:
    MvFlexpartGrid(const std::string&, const std::string&, const std::string&, const std::string&, bool forward, int species);
    ~MvFlexpartGrid() = default;

private:
    void convert();
    void convertFlux();
    void readField(MvBinaryReader*, double**, size_t, bool, double);
    void readFluxField(MvBinaryReader*, double**, size_t);
    void writeField(double* data, const std::string& outFile, long paramId, const std::string& levType, double levVal,
                    long step, int specPoint, int spec, int ageClass);
    void initData(double**, size_t);

    std::string dir_;
    MvFlexpartHeader header_;
    size_t fieldSize_;
    std::string gridFile_;
    std::string gridType_;
    std::string outPattern_;
    bool forward_{true};
    int species_{0};  // the species index, starts from 0, only used for conc!
    std::string speciesName_;
};

MvFlexpartGrid::MvFlexpartGrid(const std::string& headerFile, const std::string& gridFile,
                               const std::string& gridType, const std::string& outPattern,
                               bool forward, int species) :
    header_(headerFile),
    fieldSize_(header_.nx * header_.ny),
    gridFile_(gridFile),
    gridType_(gridType),
    outPattern_(outPattern),
    forward_(forward),
    species_(species)

{
    // feenableexcept(FE_ALL_EXCEPT & ~FE_INEXACT);

    std::cout << "Converting flexpart output file=" << gridFile << " to grib" << std::endl;
    std::cout << "species=" << species_ << std::endl;

    header_.print();

    if (gridType_ == "conc" || gridType_ == "pptv" || gridType_ == "time")
        convert();
    else if (gridType_ == "flux")
        convertFlux();
}

void MvFlexpartGrid::convert()
{
    MvBinaryReader r(gridFile_);

    if (!r.isValid()) {
        std::cerr << "Could not open flexpart output grid: " << gridFile_ << std::endl;
        exit(1);
    }

    int paramId = 0.;
    double scaleFactor = pow(10., -12);
    double* data = nullptr;
    std::string levType;
    int spec = species_;

    r.skipRecMarker();
    int step = r.read<int>();  // forecast time in seconds
    // std::cout << "TIME " << itime << std::endl;
    r.skipRecMarker();

    // bool splitByRelease = (static_cast<int>(header_.release.size()) == header_.specPointNum);

    // Loop for number of releaspoints for which a different output shall be created
    for (int specPoint = 0; specPoint < header_.specPointNum; specPoint++) {
        // Loop for age classes
        for (size_t ac = 0; ac < header_.ageClass.size(); ac++) {
            size_t dataSize = 0;

            // Wet/dry deposition
            if (forward_) {
                // Wet deposition
                levType = "sfc";
                paramId = 20198;

                dataSize = fieldSize_;
                initData(&data, dataSize);
                readField(&r, &data, dataSize, false, scaleFactor);
                writeField(data, outPattern_ + "wet.grib", paramId, levType, 0., step, specPoint, spec, ac);

                // Dry deposition
                levType = "sfc";
                paramId = 20197;

                dataSize = fieldSize_;
                initData(&data, dataSize);
                readField(&r, &data, dataSize, false, scaleFactor);
                writeField(data, outPattern_ + "dry.grib", paramId, levType, 0., step, specPoint, spec, ac);
            }
            // No depsotion for backward simulations
            else {
                initData(&data, dataSize);
                readField(&r, &data, dataSize, false, scaleFactor);
                initData(&data, dataSize);
                readField(&r, &data, dataSize, false, scaleFactor);
            }

            levType = "heightAboveGround";

            // Height level data
            if (forward_) {
                // Concentration
                if (gridType_ == "conc") {
                    // Concentration
                    if (header_.indReceptor == 1)
                        paramId = 20000;
                    // Mass mixing ratio
                    else
                        paramId = 20002;
                }
                // Volume mixing ratio
                else if (gridType_ == "pptv") {
                    paramId = 20052;
                }
                else {
                    assert(1);
                }
            }
            else {
                // Response
                if (gridType_ == "time") {
                    // Particle resident time
                    if (header_.indSource == 1 && header_.indReceptor == 1)
                        paramId = 20194;
                    // Response for mass concentration
                    else if (header_.indSource == 1 && header_.indReceptor == 2)
                        paramId = 20195;
                    // Response for mass mixing ratio
                    else if (header_.indSource == 2 && header_.indReceptor == 1)
                        paramId = 20196;
                    // Particle resident time
                    else if (header_.indSource == 2 && header_.indReceptor == 1)
                        paramId = 20194;
                    else {
                        std::cerr << "Invalid header" << std::endl;
                        exit(1);
                    }
                }
                // Volume mixing ratio
                else if (gridType_ == "pptv") {
                    paramId = 20000;
                }
            }


            dataSize = fieldSize_ * header_.levels.size();
            initData(&data, dataSize);
            readField(&r, &data, dataSize, true, scaleFactor);

            double* dLev = data;
            for (double levVal : header_.levels) {
                std::string outName = outPattern_ + std::to_string(levVal) + "_" +
                                      std::to_string(ac) + "_" +
                                      std::to_string(specPoint) + " .grib";

                writeField(dLev, outName,
                           paramId, levType, levVal, step, specPoint, spec, ac);
                dLev += fieldSize_;
            }
        }
    }

    if (data)
        delete[] data;
}

void MvFlexpartGrid::convertFlux()
{
    MvBinaryReader r(gridFile_);

    if (!r.isValid())
        return;

    double* data = nullptr;
    std::string levType = "heightAboveGround";

    r.skipRecMarker();
    int step = r.read<int>();
    r.skipRecMarker();

    size_t dataSize = fieldSize_ * header_.levels.size();

    std::vector<int> paramIds;
    paramIds.push_back(20199);  // east
    paramIds.push_back(20200);  // west
    paramIds.push_back(20201);  // south
    paramIds.push_back(20202);  // north
    paramIds.push_back(20203);  // up
    paramIds.push_back(20204);  // down

    // Loop for species
    for (int spec = 0; spec < header_.specNum; spec++) {
        // Loop for number of releaspoints for which a different output shall be created.
        // If OUTPUT_FOR_EACH_RELEASE is off.There will be just one field created merging all the
        // releases
        for (int specPoint = 0; specPoint < header_.specPointNum; specPoint++) {
            // Loop for age classes
            for (size_t ac = 0; ac < header_.ageClass.size(); ac++) {
                // Flux types
                for (size_t iFlux = 0; iFlux < paramIds.size(); iFlux++) {
                    initData(&data, dataSize);

                    // read data for all the levels
                    readFluxField(&r, &data, dataSize);

                    double* dLev = data;
                    for (double levVal : header_.levels) {
                        std::string outName = outPattern_ + std::to_string(iFlux) + "_" +
                                              std::to_string(levVal) + "_" +
                                              std::to_string(ac) + "_" +
                                              std::to_string(specPoint) + "_" +
                                              std::to_string(spec + 1) + "_.grib";

                        writeField(dLev, outName,
                                   paramIds[iFlux], levType, levVal, step, specPoint, spec, 0);
                        dLev += fieldSize_;
                    }
                }
            }
        }
    }

    if (data)
        delete[] data;
}

void MvFlexpartGrid::initData(double** data, size_t num)
{
    if (*data)
        delete[] * data;

    *data = new double[num];

    for (size_t i = 0; i < num; i++)
        (*data)[i] = 0.;
}

void MvFlexpartGrid::readField(MvBinaryReader* r, double** data, size_t dataSize, bool doAbs, double scaleFactor)
{
    int factor = 1;
    int cnt = 0;
    size_t n = 0;

    r->skipRecMarker();
    int spCnt_i = r->read<int>();
    r->skipRecMarker();

    r->skipRecMarker();
    std::vector<int> spDump_i;
    r->read<int>(spCnt_i, spDump_i);
    r->skipRecMarker();

    r->skipRecMarker();
    int spCnt_r = r->read<int>();
    r->skipRecMarker();

    r->skipRecMarker();
    std::vector<float> spDump_r;
    r->read<float>(spCnt_r, spDump_r);
    r->skipRecMarker();

    // std::cout << "spDump_i" << std::endl;
    // for(size_t i=0; i < spDump_i.size(); i++)
    //     std::cout << "  " << spDump_i[i] << std::endl;

    // std::cout << "spDump_r" << std::endl;
    // for(size_t i=0; i < spDump_r.size(); i++)
    //     std::cout << "  " << spDump_r[i] << std::endl;


    assert(spCnt_r == static_cast<int>(spDump_r.size()));

    double dVal = 0.;

    for (int i = 0; i < spCnt_r; i++) {
        if (spDump_r[i] * factor > 0.) {
            n = spDump_i[cnt];
            cnt++;
            factor *= -1;
        }
        else {
            n++;
        }

        // int yPos=n/header.nx;
        // int xPos=n-header.nx*yPos;
        int np = n;
        if (n >= fieldSize_) {
            np -= fieldSize_;
        }

        // std::cout << "np " << np << " " << i << " "  << n << " " << fieldSize_ << std::endl;
        assert(np < static_cast<int>(dataSize));
        if (*data && np < static_cast<int>(dataSize)) {
            if (doAbs)
                dVal = fabs(spDump_r[i]);
            else
                dVal = spDump_r[i];

            if (fabs((*data)[np]) > 1000000000.) {
                dVal = 0.;
            }
            else {
                dVal *= scaleFactor;
            }
            (*data)[np] = dVal;
        }

        // std::cout << n << " " << data[n] << std::endl;
    }
}

void MvFlexpartGrid::readFluxField(MvBinaryReader* r, double** data, size_t dataSize)
{
    double scaleFactor = pow(10., -12);

    r->skipRecMarker();
    int type = r->read<int>();
    r->skipRecMarker();

    if (type == 1) {
        int ival = 0;
        float fval = 999.;
        while (ival != -999) {
            r->skipRecMarker();
            ival = r->read<int>();
            fval = r->read<float>();
            // std::cout << "ival=" << ival << " fval=" << fval << std::endl;
            if (ival != -999) {
                ival -= fieldSize_;
                assert(ival >= 0);
                assert(ival < static_cast<int>(dataSize));
                (*data)[ival] = fval * scaleFactor;
            }
            r->skipRecMarker();
        }
        assert(fabs(fval - 999.) < 0.000001);
    }
    else if (type == 2) {
        for (size_t k = 0; k < header_.levels.size(); k++)
            for (int i = 0; i < header_.nx; i++) {
                std::vector<float> fvec(header_.ny);

                r->skipRecMarker();
                r->read<float>(header_.ny, fvec);
                r->skipRecMarker();
                assert(static_cast<int>(fvec.size()) == header_.ny);

                for (size_t j = 0; j < fvec.size(); j++) {
                    int pos = i + j * header_.nx + k * fieldSize_;
                    assert(pos < static_cast<int>(dataSize));
                    (*data)[pos] = fvec[j] * scaleFactor;
                    if (fabs((*data)[pos]) > 1000000000.)
                        (*data)[pos] = 0.;
                }
            }
    }
}


void MvFlexpartGrid::writeField(double* data, const std::string& outFile, long paramId,
                                const std::string& levType, double levVal, long stepInSec, int specPoint, int spec, int ageClass)
{
    grib_handle* gh = nullptr;

    size_t size = header_.nx * header_.ny;
    size_t len = 0;

    std::string sampleName("flexpart_1_grib2");
    if ((gh = grib_handle_new_from_samples(nullptr, sampleName.c_str())) == nullptr) {
        std::cerr << "Could not create grib from sample: " << sampleName << std::endl;
        exit(1);
        return;
    }

    // Pouplate Section 2
    GRIB_CHECK(grib_set_long(gh, "flexpartVersion", 902), nullptr);
    GRIB_CHECK(grib_set_long(gh, "speciesId", spec + 1), nullptr);

    // Species
    assert(spec >= 0 && spec < header_.specNum);
    std::string speciesName = header_.specNames[spec];
    len = speciesName.size();
    GRIB_CHECK(grib_set_string(gh, "speciesName", speciesName.c_str(), &len), nullptr);

    // Releases

    // One field per release
    if (static_cast<int>(header_.release.size()) == header_.specPointNum) {
        int release = specPoint;
        assert(release >= 0 && release < static_cast<int>(header_.release.size()));
        GRIB_CHECK(grib_set_long(gh, "numberOfReleases", header_.release.size()), nullptr);
        GRIB_CHECK(grib_set_long(gh, "releaseNumber", release + 1), nullptr);

        std::string releaseName = header_.release[release].name;
        len = releaseName.size();
        GRIB_CHECK(grib_set_string(gh, "releaseName", releaseName.c_str(), &len), nullptr);
    }
    // releases are merged into a field
    else {
        GRIB_CHECK(grib_set_long(gh, "numberOfReleases", header_.release.size()), nullptr);
        GRIB_CHECK(grib_set_long(gh, "releaseNumber", 0), nullptr);
    }

    // Ageclasses
    GRIB_CHECK(grib_set_long(gh, "numberOfAgeClasses", header_.ageClass.size()), nullptr);
    GRIB_CHECK(grib_set_long(gh, "ageClass", ageClass + 1), nullptr);

    long ageClassBegin = 0;
    long ageClassEnd = header_.ageClass[ageClass];

    if (ageClass > 0) {
        ageClassBegin = header_.ageClass[ageClass - 1];
    }

    GRIB_CHECK(grib_set_long(gh, "ageClassEnd", ageClassBegin), nullptr);
    GRIB_CHECK(grib_set_long(gh, "ageClassEnd", ageClassEnd), nullptr);

    // Grid
    GRIB_CHECK(grib_set_double(gh, "longitudeOfFirstGridPointInDegrees", header_.west), nullptr);
    GRIB_CHECK(grib_set_double(gh, "longitudeOfLastGridPointInDegrees", header_.east), nullptr);
    GRIB_CHECK(grib_set_double(gh, "latitudeOfFirstGridPointInDegrees", header_.south), nullptr);
    GRIB_CHECK(grib_set_double(gh, "latitudeOfLastGridPointInDegrees", header_.north), nullptr);
    GRIB_CHECK(grib_set_long(gh, "Ni", header_.nx), nullptr);
    GRIB_CHECK(grib_set_long(gh, "Nj", header_.ny), nullptr);
    GRIB_CHECK(grib_set_double(gh, "iDirectionIncrementInDegrees", header_.dx), nullptr);
    GRIB_CHECK(grib_set_double(gh, "jDirectionIncrementInDegrees", header_.dy), nullptr);
    GRIB_CHECK(grib_set_long(gh, "jScansPositively", 1), nullptr);

    // Level
    len = levType.size();
    GRIB_CHECK(grib_set_string(gh, "typeOfLevel", levType.c_str(), &len), nullptr);
    GRIB_CHECK(grib_set_double(gh, "level", levVal), nullptr);

    // Param
    GRIB_CHECK(grib_set_long(gh, "paramId", paramId), nullptr);

    // Date
    GRIB_CHECK(grib_set_long(gh, "dataDate", header_.runDate), nullptr);
    GRIB_CHECK(grib_set_long(gh, "dataTime", header_.runTime / 100), nullptr);

    if (stepInSec % 3600 == 0)
        GRIB_CHECK(grib_set_long(gh, "step", stepInSec / 3600), nullptr);
    else if (stepInSec % 60 == 0) {
        std::string stepUnits("m");
        len = stepUnits.size();
        GRIB_CHECK(grib_set_string(gh, "stepUnits", stepUnits.c_str(), &len), nullptr);
        GRIB_CHECK(grib_set_long(gh, "step", stepInSec / 60), nullptr);
    }
    else {
        std::string stepUnits("s");
        len = stepUnits.size();
        GRIB_CHECK(grib_set_string(gh, "stepUnits", stepUnits.c_str(), &len), nullptr);
        GRIB_CHECK(grib_set_long(gh, "step", stepInSec), nullptr);
    }

    GRIB_CHECK(grib_set_long(gh, "bitsPerValue", 24), nullptr);

    // for(int i=0; i < size; i++)
    //{
    //     std::cout << i << " " << data[i] << std::endl;
    // data[i]=0.;
    // }

    GRIB_CHECK(grib_set_double_array(gh, "values", data, size), nullptr);

    GRIB_CHECK(grib_write_message(gh, outFile.c_str(), "w"), nullptr);

    grib_handle_delete(gh);
}

//============================================
// Main
//============================================

int main(int argc, char** argv)
{
    if (argc != 7) {
        exit(1);
    }

    std::string headerFile(argv[1]);
    std::string gridType(argv[2]);
    std::string gridFile(argv[3]);
    std::string pattern(argv[4]);
    std::string forwardStr(argv[5]);
    std::string speciesStr(argv[6]);

    bool forward = (forwardStr == "1") ? true : false;
    int species = metview::fromString<int>(speciesStr) - 1;

    MvFlexpartGrid grid(headerFile, gridFile, gridType, pattern, forward, species);
    return 0;
}
