/// Copyright (C) 2026 Fime
// All Rights Reserved
// This code is published under the GNU Lesser General Public License (GNU LGPL)
#define BOOST_TEST_DYN_LINK
#define _USE_MATH_DEFINES
#include <math.h>
#include <functional>
#include <algorithm>
#include <array>
#include <cmath>
#include <fstream>
#include <iostream>
#include <memory>
#include <sstream>
#include <stdexcept>
#include <string>
#include <vector>
#include <boost/test/unit_test.hpp>
#include <boost/timer/timer.hpp>
#ifdef USE_MPI
#include <boost/mpi.hpp>
#endif
#include <boost/date_time/gregorian/gregorian.hpp>
#include <Eigen/Dense>
#include <StOpt/regression/LocalLinearRegression.h>
#include <StOpt/regression/LocalConstRegression.h>

#include "test/c++/tools/utils/Capa.h"
#include "test/c++/tools/utils/Tunnel.h"
#include "test/c++/tools/simulators/MeanRevertingNDSimulator.h"
#include "test/c++/tools/dp/OptimizeStorage.h"
#include "test/c++/tools/dp/DynamicProgrammingStorage.h"
#include "test/c++/tools/dp/simulatePriceControlUsingControlStorageSensi.h"
#ifdef ADEPT
#include "test/c++/tools/autodiff/calculateSensiUsingControlAdept.h"
#endif
#ifdef STANMATH
#include "test/c++/tools/autodiff/calculateSensiUsingControlStan.h"
#endif

using Eigen::ArrayXd;
using Eigen::ArrayXi;
using Eigen::ArrayXXd;
using Eigen::MatrixXd;
using std::array;
using std::cout;
using std::endl;
using std::shared_ptr;
using std::make_shared;
using std::string;
using std::vector;

namespace bg = boost::gregorian;


#if defined   __linux
#include <fenv.h>
#define enable_abort_on_floating_point_exception() feenableexcept(FE_DIVBYZERO | FE_INVALID)
#endif

// ================================================================
// Helpers dates
// ================================================================

bg::date parseDateYMDSlash(const std::string &s)
{
    // format "YYYY/MM/DD"
    int y, m, d;
    char c1, c2;
    std::istringstream iss(s);
    iss >> y >> c1 >> m >> c2 >> d;
    if (!iss || c1 != '/' || c2 != '/')
        throw std::runtime_error("Invalid date format: " + s);

    return bg::date(y, m, d);
}

bg::date parseDateYMDDash(const std::string &s)
{
    // format "YYYY-MM-DD"
    int y, m, d;
    char c1, c2;
    std::istringstream iss(s);
    iss >> y >> c1 >> m >> c2 >> d;
    if (!iss || c1 != '-' || c2 != '-')
        throw std::runtime_error("Invalid date format: " + s);

    return bg::date(y, m, d);
}

int daysBetween(const bg::date &d0, const bg::date &d1)
{
    return static_cast<int>((d1 - d0).days());
}

// ================================================================
// Gates
// ================================================================

ArrayXd porteMin(
    int date,
    double QminAtDate,
    ArrayXd QMin,
    const ArrayXd &Qgrid,
    const ArrayXd &MaxInjec,
    const ArrayXd &FactorInj)
{
    QMin(date) = QminAtDate;
    while (date > 0)
    {
        date -= 1;

        double newQmin = QMin(date);

        for (int i = 0; i < Qgrid.size() - 1; ++i)
        {
            const double left0  = Qgrid(i)     + MaxInjec(i)     * FactorInj(date);
            const double left1  = Qgrid(i + 1) + MaxInjec(i + 1) * FactorInj(date);
            const double target = QMin(date + 1);
            if (((target >= left0) && (target <= left1)) ||
                    ((target <= left0) && (target >= left1)))
            {
                const double numerator =
                    target * (Qgrid(i + 1) - Qgrid(i))
                    - (MaxInjec(i) * FactorInj(date) * Qgrid(i + 1)
                       - MaxInjec(i + 1) * FactorInj(date) * Qgrid(i));

                const double denominator =
                    (Qgrid(i + 1) - Qgrid(i))
                    + (MaxInjec(i + 1) * FactorInj(date)
                       - MaxInjec(i) * FactorInj(date));

                const double QminPt = numerator / denominator;
                newQmin = std::max(newQmin, std::max(0.0, QminPt));
            }
            else if (target <= Qgrid(0)     + MaxInjec(0)     * FactorInj(date))
            {
                newQmin = 0.0;
            }
        }

        QMin(date) = newQmin;
        if (newQmin <= 1e-10)
        {
            break;
        }
    }
    return QMin;
}

ArrayXd porteMax(
    int date,
    double QmaxAtDate,
    ArrayXd QMax,
    const ArrayXd &Qgrid,
    const ArrayXd &MaxWith,
    const ArrayXd &FactorWith,
    double QQMax)
{
    QMax(date) = QmaxAtDate;

    while (date > 0)
    {
        date -= 1;

        double newQmax = QMax(date);

        for (int i = 0; i < Qgrid.size() - 1; ++i)
        {
            const double left0  = Qgrid(i)     - MaxWith(i)     * FactorWith(date);
            const double left1  = Qgrid(i + 1) - MaxWith(i + 1) * FactorWith(date);
            const double target = QMax(date + 1);

            if (((target >= left0) && (target <= left1)) ||
                    ((target <= left0) && (target >= left1)))
            {
                const double numerator =
                    target * (Qgrid(i + 1) - Qgrid(i))
                    + (MaxWith(i) * FactorWith(date) * Qgrid(i + 1)
                       - MaxWith(i + 1) * FactorWith(date) * Qgrid(i));

                const double denominator =
                    (Qgrid(i + 1) - Qgrid(i))
                    - (MaxWith(i + 1) * FactorWith(date)
                       - MaxWith(i) * FactorWith(date));

                const double QmaxPt = numerator / denominator;

                newQmax = std::min(newQmax, std::min(QQMax, QmaxPt));
            }
            else if ((Qgrid(Qgrid.size() - 1)
                      - MaxWith(MaxWith.size() - 1) * FactorWith(date)) <= target)
            {
                newQmax = QQMax;
            }
        }

        QMax(date) = newQmax;

        if (newQmax >= QQMax - 1e-10)
            break;
    }
    return QMax;
}

// ================================================================
// Storage desciption
// ================================================================

struct StorageDescription
{
    double QMax;
    double CI;
    double CW;

    ArrayXd QGridRateI;
    ArrayXd reductFactI;

    ArrayXd QGridRateW;
    ArrayXd reductFactW;

    std::vector<std::array<double, 3>> porteValues;
    // Ici on stocke date sous forme MMDD encodée séparément.
    std::vector<std::string> porteDates;

    std::vector<std::pair<std::string, double>> factorMulInj;
    std::vector<std::pair<std::string, double>> factorMulWith;
};

ArrayXd makeArray(std::initializer_list<double> values)
{
    ArrayXd a(values.size());
    int i = 0;
    for (double v : values)
        a(i++) = v;
    return a;
}

StorageDescription makeDynStor()
{
    StorageDescription s;

    s.QMax = 1000000.0;
    const double nbJIn = 110.0;
    const double nbJW  = 59.0;

    s.CI = s.QMax / nbJIn;
    s.CW = s.QMax / nbJW;

    s.QGridRateI = makeArray({0.0, 0.6, 1.0});
    s.reductFactI = makeArray({1.0, 1.0, 1.0});

    s.QGridRateW = makeArray({0.0, 0.30778, 0.39572, 0.87938, 1.0});
    s.reductFactW = makeArray({0.21, 0.87, 1.0, 1.0, 1.0});

    s.porteDates = {"08/01", "11/01", "01/01", "03/01", "04/01"};
    s.porteValues =
    {
        {0.0, 0.3,  0.75},
        {0.0, 0.85, 1.0},
        {0.0, 0.45, 0.9},
        {0.0, 0.1,  0.55},
        {0.0, 0.0,  0.4}
    };

    s.factorMulInj  = {{"04/01", 0.817}, {"11/01", 0.993}};
    s.factorMulWith = {{"04/01", 0.994}, {"11/01", 0.993}};

    return s;
}


// ================================================================
// decodeMulFacCap
// ================================================================

ArrayXd decodeMulFacCap(
    const std::string &debYearStr,
    const std::string &endYearStr,
    const std::vector<std::pair<std::string, double>> &factors)
{
    const bg::date debDate = parseDateYMDSlash(debYearStr);
    const bg::date endDate = parseDateYMDSlash(endYearStr);

    const int nbDays = daysBetween(debDate, endDate) + 1;

    ArrayXd ret(nbDays);

    std::vector<std::pair<bg::date, double>> datedFactors;

    for (int year = debDate.year() - 1; year <= debDate.year() + 2; ++year)
    {
        for (const auto &f : factors)
        {
            const std::string &md = f.first;
            const double value = f.second;

            int m, d;
            char c;
            std::istringstream iss(md);
            iss >> m >> c >> d;

            datedFactors.push_back({bg::date(year, m, d), value});
        }
    }

    std::sort(
        datedFactors.begin(),
        datedFactors.end(),
        [](const auto & a, const auto & b)
    {
        return a.first < b.first;
    }
    );

    for (int i = 0; i < nbDays; ++i)
    {
        const bg::date cur = debDate + bg::days(i);

        double value = datedFactors.front().second;

        for (const auto &df : datedFactors)
        {
            if (df.first <= cur)
                value = df.second;
            else
                break;
        }

        ret(i) = value;
    }

    return ret;
}

//  build a future curve
Eigen::ArrayXd buildFutureCurve()
{
    std::vector<double> values;
    values.reserve(365);

    auto addRepeated = [&](double price, int n)
    {
        for (int i = 0; i < n; ++i)
            values.push_back(price);
    };

    // Avril 2024 : 30 jours
    values.push_back(26.38);       // 2024-04-01
    values.push_back(26.70);       // 2024-04-02
    addRepeated(27.55, 3);         // 2024-04-03 -> 2024-04-05
    addRepeated(27.06, 10);        // 2024-04-06 -> 2024-04-15
    addRepeated(27.15, 15);        // 2024-04-16 -> 2024-04-30

    // Mois suivants
    addRepeated(27.11, 31);        // Mai 2024
    addRepeated(27.20, 30);        // Juin 2024
    addRepeated(27.38, 31);        // Juillet 2024
    addRepeated(27.53, 31);        // Août 2024
    addRepeated(27.96, 30);        // Septembre 2024
    addRepeated(28.06, 31);        // Octobre 2024
    addRepeated(30.57, 30);        // Novembre 2024
    addRepeated(31.69, 31);        // Décembre 2024
    addRepeated(31.99, 31);        // Janvier 2025
    addRepeated(32.06, 28);        // Février 2025
    addRepeated(31.50, 31);        // Mars 2025

    if (values.size() != 365)
        throw std::runtime_error("Invalid future curve size");

    Eigen::ArrayXd fut(values.size());

    for (int i = 0; i < static_cast<int>(values.size()); ++i)
        fut(i) = values[i];

    return fut;
}


BOOST_AUTO_TEST_CASE(testStorageWithAutoDiff)
{
#ifdef USE_MPI
    boost::mpi::communicator world;
#endif

    // Paramètres équivalents aux defaults Python
    std::string debYear = "2024/04/01";
    std::string endYear = "2025/04/01";

    int nbSimOpt = 100; //00;
    int nbSimSim = 100; //00;
    int batchSize = 500;

    double stockDis = 1000.0;
    double stockInit = 0.0;

    double sigS = 0.7877;
    double aS   = 1.9641;
    double sigL = 0.4968;

    ArrayXi nbMesh(2);
    nbMesh << 1, 1;
    //nbMesh << 3, 2;

    int regType = 0;

    const bg::date debDate = parseDateYMDSlash(debYear);
    const bg::date endDate = parseDateYMDSlash(endYear);

    const int nbStep = daysBetween(debDate, endDate);

    cout << "Beginning of Year " << debYear << endl;
    cout << "End of year " << endYear << endl;
    cout << "NB days " << nbStep << endl;

    // Forward
    ArrayXd fut = buildFutureCurve();
    cout << " futSIZE " << fut.size() << " NBT " << nbStep << endl ;

    // Paramètres modèle en jours
    sigS /= std::sqrt(365.0);
    aS   /= 365.0;
    sigL /= std::sqrt(365.0);

    ArrayXd sig(2);
    sig << sigS, sigL;

    ArrayXd mr(2);
    mr << aS, 1e-7;

    MatrixXd correl = MatrixXd::Identity(2, 2);

    // Storage
    StorageDescription storData = makeDynStor();

    ArrayXd factorInj =
        decodeMulFacCap(debYear, endYear, storData.factorMulInj);

    ArrayXd factorWith =
        decodeMulFacCap(debYear, endYear, storData.factorMulWith);

    ArrayXd QMin = ArrayXd::Zero(nbStep + 1);
    ArrayXd QMax = ArrayXd::Constant(nbStep + 1, storData.QMax);

    // Portes
    const int startYear = debDate.year();

    for (size_t k = 0; k < storData.porteDates.size(); ++k)
    {
        const std::string &md = storData.porteDates[k];

        int month, day;
        char c;
        std::istringstream iss(md);
        iss >> month >> c >> day;

        int yearN;

        if ((month > 4) || ((month == 4) && (day > 1)))
            yearN = startYear;
        else
            yearN = startYear + 1;

        bg::date gateDate(yearN, month, day);
        int date = daysBetween(debDate, gateDate);

        if (date <= nbStep && date >= 0)
        {
            const double portMin = storData.porteValues[k][1] * storData.QMax;
            const double portMax = storData.porteValues[k][2] * storData.QMax;

            ArrayXd QIgrid = storData.QGridRateI * storData.QMax;
            ArrayXd VIgrid = storData.reductFactI * storData.CI;

            ArrayXd QWgrid = storData.QGridRateW * storData.QMax;
            ArrayXd VWgrid = storData.reductFactW * storData.CW;

            QMin = porteMin(
                       date,
                       portMin,
                       QMin,
                       QIgrid,
                       VIgrid,
                       factorInj
                   );

            QMax = porteMax(
                       date,
                       portMax,
                       QMax,
                       QWgrid,
                       VWgrid,
                       factorWith,
                       storData.QMax
                   );
        }
    }

    // Tunnel
    std::vector<std::array<double, 2>> tunnelV(nbStep + 1);

    for (int i = 0; i <= nbStep; ++i)
    {
        tunnelV[i] = {QMin(i), QMax(i)};
    }

    Tunnel tunnel(tunnelV, stockDis);

    // Capacity
    ArrayXd QIgrid = storData.QGridRateI * storData.QMax;
    ArrayXd VIgrid = storData.reductFactI * storData.CI;
    Capa theInjCapacity(factorInj, QIgrid, VIgrid);

    ArrayXd QWgrid = storData.QGridRateW * storData.QMax;
    ArrayXd VWgrid = storData.reductFactW * storData.CW;

    Capa theWithdrawalCapacity(factorWith, QWgrid, VWgrid);

    ArrayXd initStorageVolume(1);
    initStorageVolume << stockInit;

    // Regressor
    std::shared_ptr<StOpt::BaseRegression> regressor;

    if (regType == 0)
    {
        cout << "LINEAR PER MESH REGRESSOR" << endl;
        regressor = std::make_shared<StOpt::LocalLinearRegression>(nbMesh);
    }
    else if (regType == 1)
    {
        cout << "CONSTANT PER MESH REGRESSION" << endl;
        regressor = std::make_shared<StOpt::LocalConstRegression>(nbMesh);
    }
    else
    {
        throw std::runtime_error("Only regType 0/1 implemented in this C++ sample");
    }

    // File names
    std::string toAdd =
        std::string("Dynamic")
        + "nbSim" + std::to_string(nbSimOpt)
        + "nbMesh" + std::to_string(nbMesh(0))
        + "stDis" + std::to_string(static_cast<int>(stockDis));

    std::string baseNameForSim = "SIMUSTORE";
    std::string baseBellmanAndCont = toAdd + "CaseMeanReverting2D";

    // Optimisation backward
    bool bForward = false;
    int seedOpt = 123456;

    auto simulatorOpt =
        std::make_shared<MeanRevertingNDSimulator>(
            fut,
            sig,
            mr,
            correl,
            nbStep,
            nbSimOpt,
            bForward,
            baseNameForSim,
#ifdef USE_MPI
            world,
#endif
            seedOpt
        );

    auto optimizer =
        std::make_shared<OptimizeStorage<MeanRevertingNDSimulator>>(
            theInjCapacity,
            theWithdrawalCapacity
        );

    optimizer->setSimulator(simulatorOpt);

    DynamicProgrammingStorage(
        tunnel,
        optimizer,
        regressor,
        initStorageVolume,
        baseBellmanAndCont
#ifdef USE_MPI
        , world
#endif
    );

    // forward
    bForward = true;
    // // Simulation control + delta classique
    int seedControl = 7891;

    auto simulatorSimControl =
        std::make_shared<MeanRevertingNDSimulator>(
            fut,
            sig,
            mr,
            correl,
            nbStep,
            nbSimSim,
            bForward,
            baseNameForSim,
#ifdef USE_MPI
            world,
#endif
            seedControl
        );

    optimizer->setSimulator(simulatorSimControl);

    auto resControl =
        simulatePriceControlUsingControlStorageSensi(
            tunnel,
            optimizer,
            initStorageVolume,
            baseBellmanAndCont
#ifdef USE_MPI
            , world
#endif
        );
    std::cout << " Value with control " << std::get<2>(resControl) << std::endl ;
    std::cout << " Tangent"  << std::get<3>(resControl).transpose() << std::endl;
    std::cout << " Vol deriv " << std::get<4>(resControl).transpose() << std::endl;

    ///  Autodiff Adept
#ifdef ADEPT

    auto sensA =
        calculateSensiUsingControlAdept(
            fut,
            sig,
            mr,
            correl,
            nbStep,
            nbSimSim,
            batchSize,
            baseNameForSim,
            seedControl,
            theInjCapacity,
            theWithdrawalCapacity,
            tunnel,
            initStorageVolume,
            baseBellmanAndCont
        );
    cout << " ADEPT" << endl ;
    cout << "Grad with autodiff delta = " << std::get<0>(sensA).transpose() << endl;
    cout << "Grad with autodiff vega  = " <<  std::get<1>(sensA).transpose() << endl;
    cout << "Grad with autodiff mean reverting"   <<  std::get<2>(sensA).transpose() << endl;
    cout << "Grad with autodiff correl"   << std::get<3>(sensA) << endl;
#endif
    /// Autodiff stan math
#ifdef STANMATH

    auto sensB =
        calculateSensiUsingControlStan(
            fut,
            sig,
            mr,
            correl,
            nbStep,
            nbSimSim,
            batchSize,
            baseNameForSim,
            seedControl,
            theInjCapacity,
            theWithdrawalCapacity,
            tunnel,
            initStorageVolume,
            baseBellmanAndCont
        );
    cout << "STAN MATH" << endl ;
    cout << "Grad with autodiff delta = " << std::get<0>(sensB).transpose() << endl;
    cout << "Grad with autodiff vega  = " <<  std::get<1>(sensB).transpose() << endl;
    cout << "Grad with autodiff mean reverting"   <<  std::get<2>(sensB).transpose() << endl;
    cout << "Grad with autodiff correl"   << std::get<3>(sensB) << endl;
#endif

}

// (empty) Initialization function. Can't use testing tools here.
bool init_function()
{
    return true;
}

int main(int argc, char *argv[])
{
#if defined   __linux
    enable_abort_on_floating_point_exception();
#endif
#ifdef USE_MPI
    boost::mpi::environment env(argc, argv);
#endif
    return ::boost::unit_test::unit_test_main(&init_function, argc, argv);
}
