// Copyright (C) 2026 EDF
// All Rights Reserved
// This code is published under the GNU Lesser General Public License (GNU LGPL)
#undef USE_MPI
#include <cstdint>
#include <limits>
#include <thread>
#include <functional>
#include <string>
#include <stdexcept>
#include <tuple>
#ifdef _OPENMP
#include <omp.h>
#endif
#include <adept/adept.h>
#include <Eigen/Dense>
#include <boost/timer/timer.hpp>
#include <boost/chrono.hpp>
#include "test/c++/tools/simulators/MeanRevertingNDSimulator.h"
#include "test/c++/tools/dp/OptimizeStorage.h"
#include "test/c++/tools/dp/simulatePriceControlUsingControlStorageAutoDiff.h"
#include "test/c++/tools/utils/Capa.h"
#include "test/c++/tools/utils/Tunnel.h"
#include "test/c++/tools/utils/toDoubleAdept.h"
#include "test/c++/tools/utils/seedForBatch.h"

using namespace std;
using namespace Eigen;
using namespace StOpt;
using boost::timer::cpu_timer;
using adept::adouble;
using adept::Stack;


/// \brief Calculate sensibility using auto diff
/// \param p_future  future curve
/// \param p_sig     volatilities per factor
/// \param p_a       mean reverting per factor
/// \param p_correl correlation matrix
/// \param nbStep   number of resolution step
/// \param p_nbSimul  number of simulations for valuation
/// \param p_batchSize batch size used to gather simulations
/// \param p_nameArch name of base file to store backward simulations
/// \param p_seed seed used for random generator
/// \param p_injectionVolume  injection capacity
/// \param p_withDrawalVolume withdrawal capacity
/// \param p_tunnel tunnel constraining the position in volume storage
/// \param p_pointStock initial position in storage
/// \param p_fileForBellmanAndCont file name to store Bellman values and controls
std::tuple<ArrayXd, ArrayXd, ArrayXd, ArrayXXd>
calculateSensiUsingControlAdept(
    const ArrayXd &p_future,
    const ArrayXd &p_sig,
    const ArrayXd &p_a,
    const MatrixXd &p_correl,
    const int &p_nbStep,
    const size_t &p_nbSimul,
    const size_t &p_batchSize,
    const string &p_nameArch,
    int p_seed,
    const Capa &p_injectionVolume,
    const Capa &p_withDrawalVolume,
    const Tunnel &p_tunnel,
    const ArrayXd &p_pointStock,
    const string &p_fileForBellmanAndCont)
{
    using adept::adouble;

    if (p_batchSize == 0)
        throw std::runtime_error("p_batchSize must be strictly positive");

    if (p_nbSimul == 0)
        throw std::runtime_error("p_nbSimul must be strictly positive");

    cpu_timer timer1;
    timer1.start();

    const size_t nbBatch =
        (p_nbSimul + p_batchSize - 1) / p_batchSize;

    std::vector<ArrayXd> deltaBatch(
        nbBatch,
        ArrayXd::Zero(p_future.size())
    );

    std::vector<ArrayXd> vegaBatch(
        nbBatch,
        ArrayXd::Zero(p_sig.size())
    );
    std::vector<ArrayXd> mrSenBatch(
        nbBatch,
        ArrayXd::Zero(p_a.size())
    );
    std::vector<ArrayXXd> correlSenBatch(
        nbBatch,
        ArrayXXd::Zero(p_correl.rows(), p_correl.cols())
    );

    std::vector<double> valueBatch(nbBatch, 0.0);
    std::vector<size_t> nbSimulPerBatch(nbBatch, 0);

    for (size_t ibatch = 0; ibatch < nbBatch; ++ibatch)
    {
        const size_t firstSim = ibatch * p_batchSize;

        nbSimulPerBatch[ibatch] =
            std::min(p_batchSize, p_nbSimul - firstSim);
    }

    std::exception_ptr firstException = nullptr;
#ifdef _OPENMP
    #pragma omp parallel for schedule(dynamic)
#endif
    for (long long ibatchLL = 0; ibatchLL < static_cast<long long>(nbBatch); ++ibatchLL)
    {
        const size_t ibatch = static_cast<size_t>(ibatchLL);

        try
        {
            const size_t nbSimulBatch = nbSimulPerBatch[ibatch];

            if (nbSimulBatch == 0)
                continue;

            // Stack Adept
            adept::Stack stack;

            // Local to the batch
            Array<adouble, Eigen::Dynamic, 1> futAdept = p_future.cast<adouble>();
            Array<adouble, Eigen::Dynamic, 1> sigAdept = p_sig.cast<adouble>();
            Array<adouble, Dynamic, 1> mrAdept =   p_a.cast<adouble>();
            Matrix<adouble, Dynamic, Dynamic> correlAdept = p_correl.cast<adouble>();

            stack.new_recording();

            // Seed depends only  on p_seed and ibatch,
            // and independant of the thread number
            const int seedBatch =
                makeSeedForBatch(p_seed, ibatch);

            // useless as the simulator does not dump in forward mode.
            const string nameArchBatch =
                p_nameArch + "_batch_" + std::to_string(ibatch);

            std::shared_ptr<SimulatorDPBase> simulator =
                std::make_shared<MeanRevertingNDSimulatorT<adouble>>(
                    p_future,
                    p_sig,
                    p_a,
                    p_correl,
                    futAdept,
                    sigAdept,
                    mrAdept,
                    correlAdept,
                    p_nbStep,
                    nbSimulBatch,
                    true,
                    nameArchBatch,
                    seedBatch
                );

            std::shared_ptr<OptimizeStorage<MeanRevertingNDSimulatorT<adouble>>> optimizer =
                std::make_shared<OptimizeStorage<MeanRevertingNDSimulatorT<adouble>>>(
                    p_injectionVolume,
                    p_withDrawalVolume
                );

            optimizer->setSimulator(simulator);

            adouble value =
                simulatePriceControlUsingControlStorageAutoDiff<OptimizeStorage<MeanRevertingNDSimulatorT<adouble>>, adouble>(
                    p_tunnel,
                    optimizer,
                    p_pointStock,
                    p_fileForBellmanAndCont
                );

            value.set_gradient(1.0);
            stack.reverse();

            for (int i = 0; i < p_future.size(); ++i)
                deltaBatch[ibatch](i) = futAdept(i).get_gradient();

            for (int i = 0; i < p_sig.size(); ++i)
                vegaBatch[ibatch](i) = sigAdept(i).get_gradient();
            for (int i = 0; i < p_a.size(); ++i)
                mrSenBatch[ibatch](i) = mrAdept(i).get_gradient();
            for (int j = 0; j < p_correl.cols(); ++j)
                for (int i = 0; i < p_correl.rows(); ++i)
                    correlSenBatch[ibatch](i, j) = correlAdept(i, j).get_gradient();


            valueBatch[ibatch] = value.value();

        }
        catch (...)
        {
#ifdef _OPENMP
            #pragma omp critical
#endif
            {
                if (!firstException)
                    firstException = std::current_exception();
            }
        }
    }

    if (firstException)
        std::rethrow_exception(firstException);

    ArrayXd delta = ArrayXd::Zero(p_future.size());
    ArrayXd vega  = ArrayXd::Zero(p_sig.size());
    ArrayXd mrSensi  = ArrayXd::Zero(p_a.size());
    ArrayXXd correlSensi  = ArrayXXd::Zero(p_correl.rows(), p_correl.cols());

    double valueMean = 0.0;

    for (size_t ibatch = 0; ibatch < nbBatch; ++ibatch)
    {
        const double weight =
            static_cast<double>(nbSimulPerBatch[ibatch]) /
            static_cast<double>(p_nbSimul);

        delta += weight * deltaBatch[ibatch];
        vega  += weight * vegaBatch[ibatch];
        mrSensi += weight * mrSenBatch[ibatch];
        correlSensi += weight * correlSenBatch[ibatch];

        valueMean += weight * valueBatch[ibatch];
    }

    timer1.stop();

    boost::chrono::duration<double> durOpt1 =
        boost::chrono::nanoseconds(timer1.elapsed().user);

    const double timeSimul = durOpt1.count();

#ifdef _OPENMP
    const int nbThreads = omp_get_max_threads();
#else
    const int nbThreads = 1;
#endif

    cout << "AutoDiff Adept OpenMP batched value " << valueMean
         << " time " << timeSimul
         << " nbBatch " << nbBatch
         << " batchSize " << p_batchSize
         << " nbSimul " << p_nbSimul
         << " nbOpenMPThreads " << nbThreads
         << endl;

    return make_tuple(delta, vega, mrSensi, correlSensi);
}
