// Copyright (C) 2025 EDF
// All Rights Reserved
// This code is published under the GNU Lesser General Public License (GNU LGPL)
#include <fstream>
#include <memory>
#include <functional>
#ifdef USE_MPI
#include <boost/mpi.hpp>
#endif
#include <boost/lexical_cast.hpp>
#include <Eigen/Dense>
#include "geners/BinaryFileArchive.hh"
#include "StOpt/regression/BaseRegression.h"
#include "StOpt/core/grids/GridAdapt2D.h"
#include "StOpt/dp/FinalStepDPCut.h"
#include "StOpt/dp/TransitionStepRegressionDPCutGridAdapt.h"
#include "StOpt/dp/OptimizerDPCutGridAdaptBase.h"

using namespace std;
using namespace Eigen;


double  DynamicProgrammingByRegressionCutGridAdapt2D(vector<shared_ptr<StOpt::GridAdaptBase> > &p_grids,
        const shared_ptr<StOpt::OptimizerDPCutGridAdaptBase > &p_optimize,
        const shared_ptr<StOpt::BaseRegression> &p_regressor,
        const string   &p_fileToDump
#ifdef USE_MPI
        , const boost::mpi::communicator &p_world
#endif
                                                    )
{
    // from the optimizer get back the simulator
    shared_ptr< StOpt::SimulatorDPBase> simulator = p_optimize->getSimulator();
    // final cut values
    shared_ptr<ArrayXXd>   valueCutsNext = make_shared<ArrayXXd>(ArrayXXd::Zero(3 * simulator->getNbSimul(), p_grids[simulator->getNbStep()]->getNbPoints()));
    shared_ptr<gs::BinaryFileArchive> ar = make_shared<gs::BinaryFileArchive>(p_fileToDump.c_str(), "w");
    // name for object in archive
    string nameAr = "Continuation";

    // iterate on time steps
    for (int iStep = 0; iStep < simulator->getNbStep(); ++iStep)
    {
#ifdef USE_MPI
        if (p_world.rank() == 0)
        {
            cout << " iStep " << iStep << endl ;
        }
#else
        cout << " iStep " << iStep << endl ;
#endif
        ArrayXXd asset = simulator->stepBackwardAndGetParticles();
        // conditional expectation operator
        p_regressor->updateSimulations(((iStep == (simulator->getNbStep() - 1)) ? true : false), asset);
        // transition object
        StOpt::TransitionStepRegressionDPCutGridAdapt  transStep(p_grids[ simulator->getNbStep() - iStep - 1], p_grids[ simulator->getNbStep() - iStep], p_optimize
#ifdef USE_MPI
                , p_world
#endif
                                                                );
        shared_ptr<ArrayXXd>  valueCuts = make_shared<ArrayXXd>(transStep.oneStep(*valueCutsNext, p_regressor));
        // dump continuation values
        transStep.dumpContinuationCutsValues(ar, nameAr, iStep, *valueCutsNext,  p_regressor);

        valueCutsNext = valueCuts;
    }
    // Only keep the first nbSimul values  of the cuts  (other components are  for derivatives with respect to stock value)
    return valueCutsNext->col(0).head(simulator->getNbSimul()).mean();
}
