
/***************************** 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 "MvBinaryReader.h"

class MvFlexpartRelease
{
public:
    int num{0};
    int start{0};
    int end{0};
    std::string type;
    float west{0.};
    float east{0.};
    float north{0.};
    float south{0.};
    float top{0.};
    float bottom{0.};
    std::vector<float> massForSpec;
};


class MvFlexpartHeader
{
public:
    MvFlexpartHeader(const std::string&);

    int runDate{0};
    int runTime{0};
    int interval{0};
    int averagingInterval{0};
    int samplingInterval{0};
    float west{0.};
    float south{0.};
    float north{0.};
    float east{0.};
    int nx{0};
    int ny{0};
    float dx{0.};
    float dy{0.};
    std::vector<float> levels;
    int specNum{0};
    int specPointNum{0};
    std::vector<MvFlexpartRelease> release;
    int method{0};
    bool hasSubGrid{false};
    bool hasConvection{false};
    int indSource{0};
    int indReceptor{0};
    std::vector<int> ageClass;
};

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

protected:
    void convert();
    void readField(MvBinaryReader*, double**, size_t, bool);
    void writeField(double* data, const std::string& outFile);
    void initData(double**, size_t);

    std::string dir_;
    MvFlexpartHeader header_;
    size_t fieldSize_{0};
    std::string gridFile_;
    std::string outPattern_;
};


MvFlexpartHeader::MvFlexpartHeader(const std::string& fileName)
{
    MvBinaryReader r(fileName);

    if (!r.isValid())
        return;

    r.skipRecMarker();
    runDate = r.read<int>();
    runTime = r.read<int>();
    r.skip(13);
    r.skipRecMarker();

    r.skipRecMarker();
    interval = r.read<int>();
    averagingInterval = r.read<int>();
    samplingInterval = r.read<int>();
    r.skipRecMarker();

    r.skipRecMarker();
    west = r.read<float>();
    south = r.read<float>();
    nx = r.read<int>();
    ny = r.read<int>();
    dx = r.read<float>();
    dy = r.read<float>();
    north = south + static_cast<float>(ny) * dy;
    east = west + static_cast<float>(nx) * dx;
    r.skipRecMarker();

    r.skipRecMarker();
    int num = r.read<int>();
    r.read(num, levels);
    r.skipRecMarker();

    r.skipRecMarker();
    r.skip(8);
    r.skipRecMarker();

    r.skipRecMarker();
    specNum = r.read<int>() / 3;
    specPointNum = r.read<int>();
    r.skipRecMarker();

    // Species
    for (int i = 0; i < specNum; i++) {
        std::string wd, dd, sp;
        r.skipRecMarker();
        r.skip(4);  // 1
        r.read(10, wd);
        r.skipRecMarker();

        r.skipRecMarker();
        r.skip(4);  // 1
        r.read(10, dd);
        r.skipRecMarker();

        r.skipRecMarker();
        r.read<int>();  // z
        r.read(10, sp);
        r.skipRecMarker();

        std::cout << "spec"
                  << " " << i << " " << wd << " " << dd << " " << sp << std::endl;
    }

    // Release points

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

    for (int i = 0; i < releaseNum; i++) {
        MvFlexpartRelease rel;

        r.skipRecMarker();
        rel.start = r.read<int>();
        rel.end = r.read<int>();
        r.read(2, rel.type);
        r.skipRecMarker();

        r.skipRecMarker();
        rel.west = r.read<float>();
        rel.south = r.read<float>();
        rel.east = r.read<float>();
        rel.north = r.read<float>();
        rel.bottom = r.read<float>();
        rel.top = r.read<float>();
        r.skipRecMarker();

        r.skipRecMarker();
        rel.num = r.read<int>();
        r.skip(4);  // 1
        r.skipRecMarker();

        r.skipRecMarker();
        r.skip(45);  // compoint
        r.skipRecMarker();

        for (int j = 0; j < specNum; j++) {
            r.skipRecMarker();
            rel.massForSpec.push_back(r.read<float>());
            r.skipRecMarker();

            r.skipRecMarker();
            r.skip(4);
            r.skipRecMarker();

            r.skipRecMarker();
            r.skip(4);
            r.skipRecMarker();
        }

        release.push_back(rel);
    }

    r.skipRecMarker();
    method = r.read<int>();
    hasSubGrid = (r.read<int>() != 0);
    hasConvection = (r.read<int>() != 0);
    indSource = r.read<int>();
    indReceptor = r.read<int>();
    r.skipRecMarker();

    r.skipRecMarker();
    int acNum = r.read<int>();
    r.read<int>(acNum, ageClass);
    // r.skipRecMarker();

    // delete [] d;

    std::cout << runDate << std::endl;
    std::cout << runTime << std::endl;
    std::cout << interval << std::endl;
    std::cout << averagingInterval << std::endl;
    std::cout << samplingInterval << std::endl;
    std::cout << west << std::endl;
    std::cout << south << std::endl;
    std::cout << nx << std::endl;
    std::cout << ny << std::endl;
    std::cout << dx << std::endl;
    std::cout << dy << std::endl;
    std::cout << num << std::endl;

    for (float level : levels)
        std::cout << "  " << level << std::endl;

    std::cout << specNum << std::endl;
    std::cout << specPointNum << std::endl;
}

MvFlexpartGrid::MvFlexpartGrid(const std::string& headerFile, const std::string& gridFile, const std::string& outPattern) :
    header_(headerFile),
    fieldSize_(header_.nx * header_.ny),
    gridFile_(gridFile),
    outPattern_(outPattern)
{
    convert();
}

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

    if (!r.isValid()) {
        return;
    }

    double* data = nullptr;

    r.skipRecMarker();
    r.read<int>();  // itime
    r.skipRecMarker();

    // Loop for age classes
    for (size_t ac = 0; ac < header_.ageClass.size(); ac++) {
        size_t dataSize = 0;

        // Wet deposition
        dataSize = fieldSize_;
        initData(&data, dataSize);
        readField(&r, &data, dataSize, false);
        writeField(data, outPattern_ + "wet.grib");

        // Dry deposition
        dataSize = fieldSize_;
        initData(&data, dataSize);
        readField(&r, &data, dataSize, false);
        writeField(data, outPattern_ + "dry.grib");

        // Concentrations
        dataSize = fieldSize_ * header_.levels.size();
        initData(&data, dataSize);
        readField(&r, &data, dataSize, true);
        writeField(data, outPattern_ + "conc.grib");
    }
}

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

    *data = new double[num];

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

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

    // Wet deposition

    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();

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

    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_;
        }

        assert(np < dataSize);

        if (doAbs)
            (*data)[np] = fabs(spDump_r[i]);
        else
            (*data)[np] = spDump_r[i];

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

void MvFlexpartGrid::writeField(double* data, const std::string& outFile)
{
    grib_handle* gh = nullptr;

    size_t size = header_.nx * header_.ny;
    int err = 0;

    std::cout << "data: " << data << std::endl;

    if ((gh = grib_handle_new_from_samples(nullptr, "regular_ll_sfc_grib2")) == nullptr) {
        // perror("ERROR: could not clone field");
        return;
    }

    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);

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

    // for(int i=130; i < 140; i++)
    //    std::cout << i << " " << data[i] << std::endl;

    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)
{
    std::string headerFile(argv[1]);
    std::string gridFile(argv[2]);
    std::string pattern(argv[3]);

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