// Copyright (C) 2026 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 <Eigen/Dense>
#include <geners/BinaryFileArchive.hh>
#include "geners/vectorIO.hh"
#include "geners/Record.hh"
#include <StOpt/core/grids/FullGrid.h>
#include <StOpt/regression/BaseRegression.h>
#include <StOpt/dp/OptimizerDPBase.h>
#include <StOpt/dp/SimulatorDPBase.h>
#include <StOpt/dp/FinalStepDP.h>
#include <StOpt/dp/TransitionStepRegressionDP.h>
#include <StOpt/regression/GridAndRegressedValueGeners.h>
#include "test/c++/tools/utils/Tunnel.h"


using namespace  std;
using namespace  Eigen;
using namespace  StOpt;



double  DynamicProgrammingStorage(const Tunnel &p_tunnel,
                                  const shared_ptr<StOpt::OptimizerDPBase > &p_optimize,
                                  const shared_ptr<StOpt::BaseRegression> &p_regressor,
                                  const ArrayXd &p_pointStockInit,
                                  const string   &p_fileToDump
#ifdef USE_MPI
                                  , const boost::mpi::communicator &p_world
#endif
                                 )
{
    class ZeroFunction
    {
    public:
        ZeroFunction() {}
        double operator()(const int &, const ArrayXd &, const ArrayXd &) const
        {
            return 0. ;
        }
    };
    // test if one file generated
    shared_ptr<gs::BinaryFileArchive> ar;
#ifdef USE_MPI
    if (p_world.rank() == 0)
#endif
        ar = make_shared<gs::BinaryFileArchive>(p_fileToDump.c_str(), "w");
    // final value for gaz in  storage is zero
    function<double(const int &, const ArrayXd &, const ArrayXd &)>  funcFinalValue = ZeroFunction();
    // from the optimizer get back the simulator
    shared_ptr< StOpt::SimulatorDPBase> simulator = p_optimize->getSimulator();
    // final values on final grid
    int iStepGrid =  static_cast<int>(simulator->getCurrentStep() + StOpt::tiny);
    shared_ptr<FullGrid> gridNextStep = p_tunnel.getGrid(iStepGrid);
    vector< shared_ptr< ArrayXXd > >  valuesNext = StOpt::FinalStepDP(gridNextStep, p_optimize->getNbRegime())(funcFinalValue, simulator->getParticles().array());
    // names for object in archive
    string nameAr = "Continuation";
    string nameArBell = "Bellman";

    for (int iStep = 0; iStep < simulator->getNbStep(); ++iStep)
    {
#ifdef USE_MPI
        if (p_world.rank() == 0)
#endif
            cout << " STEP " << iStep << " ON " << simulator->getNbStep() <<   endl ;
        // before going
        // step backward
        simulator->stepBackward();
        // regression factors
        ArrayXXd regFactors = simulator->getParticles();
        // conditional expectation operator
        p_regressor->updateSimulations(((iStep == (simulator->getNbStep() - 1)) ? true : false), regFactors);
        // new grid
        iStepGrid =  static_cast<int>(simulator->getCurrentStep() + StOpt::tiny);
        shared_ptr<FullGrid> gridCurrentStep =  p_tunnel.getGrid(iStepGrid);
        // transition object
        StOpt::TransitionStepRegressionDP transStep(gridCurrentStep, gridNextStep, static_pointer_cast<StOpt::OptimizerDPBase>(p_optimize)
#ifdef USE_MPI
                , p_world
#endif
                                                   );
        pair< vector< shared_ptr< ArrayXXd > >, vector< shared_ptr< ArrayXXd > > > valuesAndControl  = transStep.oneStep(valuesNext, p_regressor);
        // store the continuaiton values and not the Bellman values
        transStep.dumpContinuationValues(ar, nameAr, iStepGrid, valuesNext, valuesAndControl.second, p_regressor);
        valuesNext = valuesAndControl.first;

        // now store the Bellman value (valuesNext has been updated), the control is not used
        transStep.dumpBellmanValues(ar, nameArBell, iStepGrid, valuesNext,  p_regressor);

        // shift grids
        gridNextStep = gridCurrentStep;
    }

    double cashMean = (gridNextStep->createInterpolator(p_pointStockInit)->applyVec(*valuesNext[0])).mean();
#ifdef USE_MPI
    if (p_world.rank() == 0)
#endif
        cout << "Mean  port " << cashMean << endl ;
    return cashMean;
}
