// Copyright (C) 2026 EDF
// All Rights Reserved
// This code is published under the GNU Lesser General Public License (GNU LGPL)
// Developped for Senergy project
#ifndef MEANREVERTINGNDSIMULATOR_H
#define MEANREVERTINGNDSIMULATOR_H
#include <memory>
#include <random>
#ifdef USE_MPI
#include <boost/mpi.hpp>
#endif
#include <Eigen/Dense>
#include "geners/BinaryFileArchive.hh"
#include "geners/Record.hh"
#include "geners/Reference.hh"
#include "StOpt/core/utils/eigenGeners.h"
#include "StOpt/dp/SimulatorDPBase.h"
#include "test/c++/tools/utils/choleskyBasic.h"


/* \file MeanRevertingNDSimulator.h
* \brief Simulate a future deformation with
 *        a given number of uncertainties
 *        A forward and a backward mode are provided for \f$T\f$ given
 *        \f$ dF(t,T) = F(t,T) ( \sum_{i=1}^N e^{-a_i (T-t)} \sigma_i (dW^i_t + \lambda dt)) \f$
 *        Brownian are correlated
 *        Simulation at date t0 for period [Tb,Te] . If nbStepInit =0, t0=Tb
  * \author Xavier Warin
 */

class MeanRevertingNDSimulator: public StOpt::SimulatorDPBase
{
private:
    Eigen::ArrayXd m_future ; // future
    Eigen::ArrayXd m_sig    ; ///< vol of each factor
    Eigen::ArrayXd m_a      ; ///< mean reverting for each factor
    Eigen::MatrixXd m_correl ; // correlation matrix between brownian factors
    Eigen::MatrixXd m_OUProcess; ///< store the Ornstein Uhlenbeck process
    double m_trend ; // calculate trend
    Eigen::MatrixXd m_correlFac; // factorized correlations
    Eigen::ArrayXd m_expActuE;
    Eigen::ArrayXd m_etypE;  ///< Monte Carlo simulation number

protected:
    size_t  m_nbSimul ; ///< Monte Carlo simulation number
    size_t  m_nbStep ; ///< number of time steps from Tb tof Te (end of study)
    double m_step  ; // time step =1 but keep for lisibility
    std::mt19937 m_generator;
    std::normal_distribution<>  m_normalDistrib;
    int m_currentStep ; ///< current step number (0 to m_nbStep+1)  to get  the OU value
    bool m_bForward ; ///< True if simulator  forward mode
    std::shared_ptr<gs::BinaryFileArchive> m_arch ; ///<Binary archive to store simulations


    void updateTrend()
    {
        m_trend = 0;
        double curDate = getCurrentStep();
        for (int i = 0; i < m_correl.rows(); ++i)
        {
            // diag term
            m_trend += m_sig(i) * m_sig(i) * (1 - exp(-2 * m_a(i) * curDate)) / (2 * m_a(i));
            // extra diag term (*2)
            for (int  j = 0; j < i; ++j)
            {
                double aE1E2 = m_correl(i, j) * (1 - exp(- (m_a(i) + m_a(j)) * curDate)) / (m_a(i) + m_a(j));
                m_trend += 2 * m_sig(i) * m_sig(j) * aE1E2;
            }
        }
        m_trend *= -0.5;
    };



    /// simulate forward on step
    void stepForwardOU()
    {
        Eigen::VectorXd normalVar(m_correlFac.rows());
        // loop on simulations
        for (size_t is = 0; is < m_nbSimul; ++is)
        {
            for (int id = 0 ; id < m_correlFac.rows(); ++id)
            {
                normalVar(id) =  m_normalDistrib(m_generator) ; //m_normalRand();
            }
            Eigen::VectorXd normalVarCorrel = m_correlFac * normalVar;
            // ou
            for (size_t id = 0; id < size_t(m_correlFac.rows()); ++id)
            {
                m_OUProcess(id, is) = m_OUProcess(id, is) * m_expActuE(id) + m_etypE(id) * normalVarCorrel(id);
            }
        }
    }

    void updateSeed(const int &p_seed)
    {
        if (p_seed >= 1)
            m_generator = std::mt19937(p_seed);
        else
        {
            std::random_device rd;
            std::seed_seq sd{rd(), rd(), rd(), rd()}; // 4 values are enough. 8 max.
            m_generator = std::mt19937(sd); // advanced seeding
        }
    }

public:
    /// \brief Constructor
    /// \parma   p_future    future curve
    /// \param   p_sig       vol of each factor (in days)
    /// \param   p_a         mean reverting for each factor  (in days)
    /// \param   p_correl    correlation matrix between brownian factors
    /// \param   p_nbStep    number of time steps (number of days)
    /// \param   p_nbSimul   number of simulations
    /// \param   p_bForward  true if forward mode for simulator
    /// \param   p_nameArch  name of the geners archive to store simulations
    /// \param   p_world     MPI communicator
    /// \param   p_seed      seed for random generator  : default  No seed
    MeanRevertingNDSimulator(const Eigen::ArrayXd &p_future,
                             const Eigen::ArrayXd &p_sig, const Eigen::ArrayXd &p_a,
                             const Eigen::MatrixXd &p_correl,
                             const int &p_nbStep,  const size_t &p_nbSimul,
                             const bool &p_bForward, const std::string   &p_nameArch,
#ifdef USE_MPI
                             const boost::mpi::communicator &p_world,
#endif
                             int p_seed = -1):  m_future(p_future), m_sig(p_sig), m_a(p_a), m_correl(p_correl),
        m_correlFac(p_correl.rows(), p_correl.cols()), m_expActuE(p_correl.rows()),
        m_etypE(p_correl.rows()),
        m_nbSimul(p_nbSimul), m_nbStep(p_nbStep), m_step(1.),
        m_normalDistrib(0., 1.), m_bForward(p_bForward)
    {
        updateSeed(p_seed);

        // simulate and store OU processus
        m_OUProcess = Eigen::MatrixXd::Constant(p_correl.rows(), m_nbSimul, 0.);

        // initial dump
#ifdef USE_MPI
        if (p_world.rank() == 0)
#endif
        {
            m_arch = std::make_shared<gs::BinaryFileArchive>(p_nameArch.c_str(), "w");
            *m_arch << gs::Record(m_OUProcess, "allStep", "SimulOU");
        }

        // time steps are equal so factorize correlatoin matrix ones for all
        Eigen::MatrixXd correloc = Eigen::MatrixXd::Zero(p_correl.rows(), p_correl.rows());
        // correlation between brownians to correlation between factors
        for (size_t i = 0; i < size_t(p_correl.rows()); ++i)
        {
            for (size_t j = 0; j < i; ++j)
            {
                double aE1E2 = (1 - exp(- (m_a(i) + m_a(j)) * m_step)) / (m_a(i) + m_a(j));
                double aE1 = sqrt((1 - exp(-2 * m_a(i) * m_step)) / (2 * m_a(i)));
                double aE2 = sqrt((1 - exp(-2 * m_a(j) * m_step)) / (2 * m_a(j)));
                correloc(i, j) = m_correl(i, j) * aE1E2 / (aE1 * aE2);
            }
        }
        // symmetrize, add identity
        correloc = correloc + correloc.transpose() + Eigen::MatrixXd::Identity(p_correl.rows(), p_correl.rows());
        // factorize
        m_correlFac = correloc.llt().matrixL();
        // factorize calculations
        for (size_t i = 0; i < size_t(m_correlFac.rows()); ++i)
        {
            m_expActuE(i) =  exp(- m_a(i) * m_step);
            double aE = sqrt((1 - exp(-2 * m_a(i) * m_step)) / (2 * m_a(i)));
            m_etypE(i) = aE;
        }
        // loop on time steps
        for (size_t i = 0; i < m_nbStep; ++i)
        {
            stepForwardOU();
#ifdef USE_MPI
            if (p_world.rank() == 0)
#endif
                *m_arch << gs::Record(m_OUProcess, "allStep", "SimulOU");
        }
#ifdef USE_MPI
        if (p_world.rank() == 0)
#endif
            m_arch->flush();
#ifdef USE_MPI
        p_world.barrier();
#endif
        // all processor read
        m_arch = std::make_shared<gs::BinaryFileArchive>(p_nameArch.c_str(), "r");
        if (p_bForward)
        {
            updateSeed(p_seed);
            m_currentStep = 0;// a jump of 1 because of initial jump
            m_OUProcess = Eigen::MatrixXd::Zero(m_OUProcess.rows(), m_OUProcess.cols());
        }
        else
        {
            m_currentStep = m_nbStep; // a jump of 1 because of initial jump
            gs::Reference< Eigen::MatrixXd > (*m_arch, "allStep", "SimulOU").restore(m_currentStep, &m_OUProcess);
        }

        // trend
        updateTrend();
    }

    /// \brief Special function for the spot
    ///
    Eigen::ArrayXd  getSpot() const
    {
        return ((m_OUProcess.array().colwise() * m_sig).colwise().sum() + m_trend).exp().transpose() * m_future(m_currentStep);
    }
    /// \brief :from OU process reconstruct the price
    /// \param p_OU  Ornstein Uhlenbecj processes
    double  getSpotUncertainties(const Eigen::ArrayXd   &p_OU) const
    {
        return m_future(m_currentStep) * exp(m_trend + (p_OU * m_sig).sum());
    }

    Eigen::ArrayXd getTangent() const
    {
        return ((m_OUProcess.array().colwise() * m_sig).colwise().sum() + m_trend).exp().transpose();
    }

    double getTangent(const Eigen::ArrayXd   &p_OU) const
    {
        return exp(m_trend + (p_OU * m_sig).sum());
    }


    /// \brief initialize (keep same simulation)
    void resetDirection(const bool &bForward)
    {
        m_bForward = bForward;
        if (m_bForward)
        {
            m_currentStep = 0;
            m_OUProcess = Eigen::MatrixXd::Zero(m_OUProcess.rows(), m_OUProcess.cols());
        }
        else
        {
            m_currentStep = m_nbStep;
            gs::Reference< Eigen::MatrixXd > (*m_arch, "allStep", "SimulOU").restore(m_currentStep, &m_OUProcess);
        }
        // trend
        updateTrend();
    }


    /// \brief get current regression factor
    Eigen::MatrixXd getParticles() const
    {
        return m_OUProcess;
    }

    /// \brief a step forward for simulations
    void  stepForward()
    {
        m_currentStep += 1;
        stepForwardOU();
        updateTrend();
    }

    /// \return  the asset values (asset,simulations)
    void  stepBackward()
    {
        m_currentStep -= 1;
        gs::Reference< Eigen::MatrixXd > (*m_arch, "allStep", "SimulOU").restore(m_currentStep, &m_OUProcess);
        updateTrend();
    }

    /// \brief a step forward for simulations
    /// \return  the asset values (asset,simulations)
    Eigen::MatrixXd  stepForwardAndGetParticles()
    {
        m_currentStep += 1;
        stepForwardOU();
        updateTrend();
        return m_OUProcess;
    }


    /// \brief a step backward for simulations
    /// \return  the asset values (asset,simulations)
    Eigen::MatrixXd stepBackwardAndGetParticles()
    {
        m_currentStep -= 1;
        gs::Reference< Eigen::MatrixXd > (*m_arch, "allStep", "SimulOU").restore(m_currentStep, &m_OUProcess);
        updateTrend();
        return m_OUProcess;
    }
    ///@{
    inline double getStep() const
    {
        return  m_step;
    }
    inline int getNbSimul() const
    {
        return m_nbSimul;
    }
    // number of time step of resolution (not including he initail step)
    int getNbStep() const
    {
        return m_nbStep;
    }
    int getDimension() const
    {
        return m_correl.rows();
    }
    // get the current date : m_currentStep-1 is the number of time step from Tb and we add the m_nbStepInt initial step
    double getCurrentStep() const
    {
        return m_currentStep * m_step;
    }
    ///@}
    ///@{
    /// actualization with the interest rate  : 0
    inline double getActuStep() const
    {
        return  1. ;
    }
    /// actualization for initial date
    inline double getActu() const
    {
        return  1.;
    }
    ///@}
    /// derivative  of trend with respect to Short term vol
    double getTrendDeriv(const int   &p_ivol)
    {
        double curDate = getCurrentStep();
        double trendDeriv = 2 * m_sig(p_ivol) * (1 - exp(-2 * m_a(p_ivol) * curDate)) / (2 * m_a(p_ivol));

        for (int i = 0; i < m_correl.rows(); ++i)
        {
            if (i !=  p_ivol)
            {
                trendDeriv += 2 * m_sig(i) * m_correl(i, p_ivol) * (1 - exp(- (m_a(p_ivol) + m_a(i)) * curDate)) / (m_a(i) + m_a(p_ivol));
            }
        }
        return  -0.5 * trendDeriv;
    }
};


template< typename T>
class MeanRevertingNDSimulatorT: public MeanRevertingNDSimulator
{
private:
    Eigen::Array<T, Eigen::Dynamic, 1>  m_futureAD ; // future AD version
    Eigen::Array<T, Eigen::Dynamic, 1>  m_sigAD    ; ///< vol of each factor AD version
    Eigen::Array<T, Eigen::Dynamic, 1>  m_aAD      ; ///< mean reverting for each factor  AD version
    Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>  m_correlAD ; // correlation matrix between brownian factors r AD version
    Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>  m_OUProcessAD; ///< store the Ornstein Uhlenbeck process  AD version
    T  m_trendAD ; // calculate trend  AD version
    Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic> m_correlFacAD; // factorized correlations  AD version
    Eigen::Array<T, Eigen::Dynamic, 1>  m_expActuEAD;  // AD version
    Eigen::Array<T, Eigen::Dynamic, 1>  m_etypEAD; // AD version

protected:


    void updateTrendAD()
    {
        m_trendAD = 0;
        double curDate = getCurrentStep();
        for (int i = 0; i < m_correlAD.rows(); ++i)
        {
            // diag term
            m_trendAD += m_sigAD(i) * m_sigAD(i) * (1 - exp(-2 * m_aAD(i) * curDate)) / (2 * m_aAD(i));
            // extra diag term (*2)
            for (int  j = 0; j < i; ++j)
            {
                T aE1E2 = m_correlAD(i, j) * (1 - exp(- (m_aAD(i) + m_aAD(j)) * curDate)) / (m_aAD(i) + m_aAD(j));
                m_trendAD += 2 * m_sigAD(i) * m_sigAD(j) * aE1E2;
            }
        }
        m_trendAD *= -0.5;
    };

    /// simulate forward on step
    void stepForwardOUAD()
    {
        Eigen::Matrix< T,  Eigen::Dynamic, 1 >  normalVar(m_correlFacAD.rows());
        // loop on simulations
        for (size_t is = 0; is < m_nbSimul; ++is)
        {
            for (int id = 0 ; id < m_correlFacAD.rows(); ++id)
            {
                normalVar(id) =  T(m_normalDistrib(m_generator)) ; //m_normalRand();
            }
            Eigen::Matrix< T,  Eigen::Dynamic, 1 > normalVarCorrel(normalVar.size());
            for (size_t i = 0; i < m_correlFacAD.rows(); ++i)
            {
                normalVarCorrel(i) = 0;
                for (size_t j = 0; j < m_correlFacAD.cols(); ++j)
                    normalVarCorrel(i) += m_correlFacAD(i, j) * normalVar(j);
            }
            // ou
            for (size_t id = 0; id < size_t(m_correlFacAD.rows()); ++id)
            {
                m_OUProcessAD(id, is) = m_OUProcessAD(id, is) * m_expActuEAD(id) + m_etypEAD(id) * normalVarCorrel(id);
            }
        }
    }

public:
    /// \brief Constructor
    /// \parma   p_future    future curve
    /// \param   p_sig       vol of each factor (in days)
    /// \param   p_a         mean reverting for each factor  (in days)
    /// \param   p_correl    correlation matrix between brownian factors
    /// \parma   p_futureAD  future curve AD version
    /// \param   p_sigAD     vol of each factor (in days) AD version
    /// \param   p_aAD       mean reverting for each factor  (in days) AD version
    /// \param   p_correlAD  correlation matrix between brownian factors AD version
    /// \param   p_nbStep    number of time steps (number of days)
    /// \param   p_nbSimul   number of simulations
    /// \param   p_bForward  true if forward mode for simulator
    /// \param   p_nameArch  name of the geners archive to store simulations
    /// \param   p_world     MPI communicator
    /// \param   p_seed      seed for random generator  : default  No seed
    MeanRevertingNDSimulatorT(const Eigen::ArrayXd &p_future,
                              const Eigen::ArrayXd &p_sig, const Eigen::ArrayXd &p_a,
                              const Eigen::MatrixXd &p_correl,
                              const Eigen::Array<T, Eigen::Dynamic, 1> &p_futureAD,
                              const Eigen::Array<T, Eigen::Dynamic, 1> &p_sigAD, const Eigen::Array<T, Eigen::Dynamic, 1> &p_aAD,
                              const Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic> &p_correlAD,
                              const int &p_nbStep,  const size_t &p_nbSimul,
                              const bool &p_bForward, const std::string   &p_nameArch,
#ifdef USE_MPI
                              const boost::mpi::communicator &p_world,
#endif
                              int p_seed = -1):  MeanRevertingNDSimulator(p_future, p_sig,  p_a, p_correl, p_nbStep, p_nbSimul, p_bForward, p_nameArch,
#ifdef USE_MPI
                                          p_world,
#endif
                                          p_seed), m_futureAD(p_futureAD), m_sigAD(p_sigAD), m_aAD(p_aAD), m_correlAD(p_correlAD),
        m_correlFacAD(p_correl.rows(), p_correl.cols()), m_expActuEAD(p_correl.rows()),
        m_etypEAD(p_correl.rows())
    {
        updateSeed(p_seed);

        // simulate and store OU processus
        m_OUProcessAD = Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>::Constant(p_correl.rows(), m_nbSimul, T(0.));

        // time steps are equal so factorize correlatoin matrix ones for all
        Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>  correlocAD = Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>::Zero(p_correl.rows(), p_correl.rows());
        // correlation between brownians to correlation between factors
        for (size_t i = 0; i < size_t(p_correl.rows()); ++i)
        {
            for (size_t j = 0; j < i; ++j)
            {
                T aE1E2 = (1 - exp(- (m_aAD(i) + m_aAD(j)) * m_step)) / (m_aAD(i) + m_aAD(j));
                T aE1 = sqrt((1 - exp(-2 * m_aAD(i) * m_step)) / (2 * m_aAD(i)));
                T aE2 = sqrt((1 - exp(-2 * m_aAD(j) * m_step)) / (2 * m_aAD(j)));
                correlocAD(i, j) = m_correlAD(i, j) * aE1E2 / (aE1 * aE2);
            }
        }
        // symmetrize, add identity
        correlocAD = correlocAD + correlocAD.transpose() + Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>::Identity(p_correl.rows(), p_correl.rows());
        // factorize
        m_correlFacAD = choleskyLowerBasic(correlocAD);
        // factorize calculations
        for (size_t i = 0; i < size_t(m_correlFacAD.rows()); ++i)
        {
            m_expActuEAD(i) =  exp(- m_aAD(i) * m_step);
            T aE = sqrt((1 - exp(-2 * m_aAD(i) * m_step)) / (2 * m_aAD(i)));
            m_etypEAD(i) = aE;
        }
        updateTrendAD();
    }


    /// \brief Special function for the spot
    Eigen::Array<T, Eigen::Dynamic, 1>  getSpotAD() const
    {
        Eigen::Array<T, Eigen::Dynamic, 1> ret(m_OUProcessAD.cols());
        const T future = m_futureAD(m_currentStep);

        for (int j = 0; j < m_OUProcessAD.cols(); ++j)
        {
            using std::exp;

            T x =  m_trendAD;
            for (int i = 0; i < m_OUProcessAD.rows(); ++i)
                x = x + m_OUProcessAD(i, j) * m_sigAD(i);

            ret(j) = exp(x) * future;
        }

        return ret;

    }

    /// \brief :from OU process reconstruct the price
    /// \param p_OU  Ornstein Uhlenbecj processes
    T  getSpotUncertaintiesAD(const Eigen::Array<T, Eigen::Dynamic, 1>   &p_OU) const
    {
        T sumOU = 0;
        for (int i = 0; i < p_OU.size(); ++i)
            sumOU += T(p_OU(i)) * m_sigAD(i);
        return m_futureAD(m_currentStep) * exp(m_trendAD + sumOU);
    }


    /// \brief get current regression factor
    Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>  getParticlesAD() const
    {
        return m_OUProcessAD;
    }

    /// \brief a step forward for simulations
    void  stepForwardAD()
    {
        m_currentStep += 1;
        stepForwardOUAD();
        updateTrendAD();
    }


    Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>  stepForwardAndGetParticlesAD()
    {
        m_currentStep += 1;
        stepForwardOUAD();
        updateTrendAD();
        return m_OUProcessAD;
    }
};

#endif
