// 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 <stan/math.hpp> // always before Eigen
#include <Eigen/Dense>
#include <tbb/parallel_for.h>
#include <tbb/blocked_range.h>
#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/toDoubleStanMath.h"
#include "test/c++/tools/utils/Capa.h"
#include "test/c++/tools/utils/Tunnel.h"
#include "test/c++/tools/utils/seedForBatch.h"

using namespace std;
using namespace Eigen;
using namespace StOpt;
using boost::timer::cpu_timer;
using stan::math::var;


namespace
{
inline int getTbbNumThreadsFromEnv()
{
    const char *env = std::getenv("STOPT_TBB_NUM_THREADS");

    if (!env)
        return 0; // 0 => on laisse TBB décider

    try
    {
        int n = std::stoi(std::string(env));

        if (n <= 0)
            return 0;

        return n;
    }
    catch (...)
    {
        throw std::runtime_error(
            "Invalid value for STOPT_TBB_NUM_THREADS. Expected positive integer."
        );
    }
}
}


/// \brief Calculate sensibility using auto diff : TBB parallelization
/// \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 abse 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 poisiotn 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>
calculateSensiUsingControlStan(
    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)
{
    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);
    }

    auto runBatches = [&]()
    {
        tbb::parallel_for(
            tbb::blocked_range<size_t>(0, nbBatch),
            [&](const tbb::blocked_range<size_t> &range)
        {
            for (size_t ibatch = range.begin(); ibatch != range.end(); ++ibatch)
            {
                const size_t nbSimulBatch = nbSimulPerBatch[ibatch];

                if (nbSimulBatch == 0)
                    continue;

                try
                {
                    // Local stan data
                    Array<var, Dynamic, 1> futStan = p_future.cast<var>();
                    Array<var, Dynamic, 1> sigStan = p_sig.cast<var>();
                    Array<var, Dynamic, 1> mrStan =   p_a.cast<var>();
                    Matrix<var, Dynamic, Dynamic> correlStan = p_correl.cast<var>();
                    // 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);

                    shared_ptr<SimulatorDPBase> simulator =
                        make_shared<MeanRevertingNDSimulatorT<var>>(
                            p_future,
                            p_sig,
                            p_a,
                            p_correl,
                            futStan,
                            sigStan,
                            mrStan,
                            correlStan,
                            p_nbStep,
                            nbSimulBatch,
                            true,
                            nameArchBatch,
                            seedBatch
                        );

                    shared_ptr<OptimizeStorage<MeanRevertingNDSimulatorT<var>>> optimizer =
                        make_shared<OptimizeStorage<MeanRevertingNDSimulatorT<var>>>(
                            p_injectionVolume,
                            p_withDrawalVolume
                        );

                    optimizer->setSimulator(simulator);

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

                    stan::math::set_zero_all_adjoints();
                    value.grad();

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

                    for (int i = 0; i < p_sig.size(); ++i)
                        vegaBatch[ibatch](i) = sigStan(i).adj();
                    for (int i = 0; i < p_a.size(); ++i)
                        mrSenBatch[ibatch](i) = mrStan(i).adj();
                    for (int j = 0; j < p_correl.cols(); ++j)
                        for (int i = 0; i < p_correl.rows(); ++i)
                            correlSenBatch[ibatch](i, j) = correlStan(i, j).adj();


                    valueBatch[ibatch] = value.val();

                    // Nettoyage mémoire Stan local au batch / thread
                    stan::math::recover_memory();
                }
                catch (...)
                {
                    stan::math::recover_memory();
                    throw;
                }
            }
        }
        );
    };

    const int nbTbbThreads = getTbbNumThreadsFromEnv();

    if (nbTbbThreads > 0)
    {
        tbb::global_control control(
            tbb::global_control::max_allowed_parallelism,
            static_cast<std::size_t>(nbTbbThreads)
        );

        runBatches();
    }
    else
    {
        // let TBB choose the number of threads.
        runBatches();
    }

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

    cout << "AutoDiff Stan TBB batched value " << valueMean
         << " time " << timeSimul
         << " nbBatch " << nbBatch
         << " batchSize " << p_batchSize
         << " nbSimul " << p_nbSimul;

    if (nbTbbThreads > 0)
        cout << " nbTbbThreads " << nbTbbThreads;
    else
        cout << " nbTbbThreads auto";

    cout << endl;

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