// Copyright (C) 2025 EDF
// All Rights Reserved
// This code is published under the GNU Lesser General Public License (GNU LGPL)
#ifndef OPTIMIZEGASSTORAGECUTGRIDADAPT2D_H
#define OPTIMIZEGASSTORAGECUTGRIDADAPT2D_H
#include "ClpSimplex.hpp"
#include <memory>
#include <vector>
#include <utility>
#include <Eigen/Dense>
#include "StOpt/core/utils/StateWithStocks.h"
#include "StOpt/core/grids/GridAdapt2D.h"
#include "StOpt/regression/BaseRegression.h"
#include "StOpt/regression/ContinuationCutsGridAdaptNonConcave.h"
#include "StOpt/dp/OptimizerDPCutGridAdaptBase.h"

/** \file OptimizeGasStorageCutGrid2DAdapt.h
 *  \brief  Simple example of a gas storage optimizer
 *          - injection rate, withdrawal rates are independent of the storage level
 *          - the size of the storage is constant, minimum gas level is 0
 *          .
 *          Designed to work in parallel/multi threaded framework
 *          The local  optimization is solved using a LP on one time step
 *          Regression used for conditional expectation
 *          Here the grid is 2D and  adaptive and no concavity is required.
 *  \author Xavier Warin
 */

/// \class OptimizeGasStorageCutGridAdapt2D OptimizeGasStorageCutGridAdapt2D.h
/// Defines a simple gas storage in 2D for optimization and simulation using cuts and a LP solver
/// The grid is 2D and adaptative.
/// No concavity supposed
/// No constraints on the storage at the end of optimization period (so the storage will be empty)
/// - when injecting the gain is  \f$ - C_{inj} ( S+ \kappa_{inj} )\f$
/// - when withdrawing the gain is  \f$  C_{with} ( S- \kappa_{with} )\f$
/// .
template< class Simulator>
class OptimizeGasStorageCutGridAdapt2D : public StOpt::OptimizerDPCutGridAdaptBase
{
private :

    /// \brief Physical constraints
    //@{
    double m_injectionRate ; ///< injection  capacity (volume) per time step \f$ C_{inj} \f$
    double m_withdrawalRate ; ///< withdrawal rate (volume) per time step \f$ C_{with} \f$
    double m_injectionCost; ///< injection cost \f$ \kappa_{inj} \f$ per volume unit
    double m_withdrawalCost ; /// withdrawal cost  \f$ \kappa_{with} \f$ per volume unit
    //@}
    /// store the simulator
    std::shared_ptr<Simulator> m_simulator;
    bool m_bGatherConcaveGrids; // true when concave grids are gathered
    int m_iStep ; // store current time step
    bool m_bRefine; // do we refine

public :

    /// \brief Constructor
    /// \param  p_injectionRate       injection rate per time step
    /// \param  p_withdrawalRate      withdrawal rate between two time steps
    /// \param  p_injectionCost       injection cost
    /// \param  p_withdrawalCost      withdrawal cost
    /// \param  p_bGatherConcaveGrids when true concave grids are gathered to  limit number of resolutions
    /// \param  p_bRefine             if true refine all grid once
    OptimizeGasStorageCutGridAdapt2D(const double   &p_injectionRate, const double &p_withdrawalRate,
                                     const double &p_injectionCost, const double &p_withdrawalCost,  const bool &p_bGatherConcaveGrids,
                                     const bool &p_bRefine):
        m_injectionRate(p_injectionRate), m_withdrawalRate(p_withdrawalRate), m_injectionCost(p_injectionCost),
        m_withdrawalCost(p_withdrawalCost), m_bGatherConcaveGrids(p_bGatherConcaveGrids), m_iStep(10000), m_bRefine(p_bRefine)
    {}

    /// \brief Useless here
    std::vector< std::array< double, 2> > getCone(const  std::vector<  std::array< double, 2>  > &p_regionByProcessor) const
    {
        return std::vector< std::array< double, 2> > ();
    }

    /// \brief Useless here
    Eigen::Array< bool, Eigen::Dynamic, 1> getDimensionToSplit() const
    {
        return Eigen::Array< bool, Eigen::Dynamic, 1> ();
    }


    ///  \brief LP creation
    ///  \param p_cuts          array of cuts for current  for current stock points
    ///  \param xMin            minimal first stock value  imposed at end of optimization
    ///  \param xMax            maximal first stock value  imposed at end of optimization
    ///  \param yMin            minimal second stock value  imposed at end of optimization
    ///  \param yMax            maximal second stock value  imposed at end of optimization
    ///  \param p_stock         currentstock levels
    ///  \param p_spot          spot price
    ///  \param p_valueAndDerivatives  function value and derivative for current simulation
    ///  \param p_stateFollowing storage reached
    ///  \param p_gain           gain reached
    void   createAndSolveLP(const Eigen::ArrayXXd &p_cuts,
                            const double &p_xMin,
                            const double &p_xMax,
                            const double &p_yMin,
                            const double &p_yMax,
                            const Eigen::ArrayXd &p_stock,
                            const double &p_spot,
                            Eigen::ArrayXd &p_valueAndDerivatives,
                            Eigen::ArrayXd &p_stateFollowing,
                            double &p_gain) const
    {
        Eigen::ArrayXi rows(6 + 3 * p_cuts.cols()); // row position
        Eigen::ArrayXi columns(6 + 3 * p_cuts.cols()) ; // columns  position
        Eigen::ArrayXd elements(6 + 3 * p_cuts.cols()) ; // constraints matrix values
        // bounds on values
        Eigen::ArrayXd lowBound(7);
        Eigen::ArrayXd upperBound(7);

        // injection first storage
        lowBound(0) = 0;
        upperBound(0) = m_injectionRate;
        // withdrawal first storage
        lowBound(1) = - m_withdrawalRate;
        upperBound(1) = 0. ;
        // bound on first storage value at end of step
        lowBound(2) =  p_xMin;
        upperBound(2) = p_xMax;
        // injection second storage
        lowBound(3) = 0;
        upperBound(3) = m_injectionRate;
        // withdrawal second storage
        lowBound(4) = - m_withdrawalRate;
        upperBound(4) = 0. ;
        // bound on first storage value at end of step
        lowBound(5) =  p_yMin;
        upperBound(5) = p_yMax;
        // for Bellman cuts
        lowBound(6) = - StOpt::infty;
        upperBound(6) =  StOpt::infty;
        // objective function
        Eigen::ArrayXd objFunc = Eigen::ArrayXd::Zero(7);
        objFunc(0) = -(p_spot + m_injectionCost);
        objFunc(1) = -(p_spot - m_withdrawalCost);
        objFunc(3) = -(p_spot + m_injectionCost);
        objFunc(4) = -(p_spot - m_withdrawalCost);
        objFunc(6) = 1.;
        // flow constraints
        // first stock
        rows(0) = 0;
        columns(0) = 0;
        elements(0) = -1;
        rows(1) = 0;
        columns(1) = 1;
        elements(1) = -1;
        rows(2) = 0;
        columns(2) = 2;
        elements(2) = 1;
        // second stock
        rows(3) = 1;
        columns(3) = 3;
        elements(3) = -1;
        rows(4) = 1;
        columns(4) = 4;
        elements(4) = -1;
        rows(5) = 1;
        columns(5) = 5;
        elements(5) = 1;
        // bounds associated to matrix constraints
        Eigen::ArrayXd lowBoundConst(2 + p_cuts.cols());
        Eigen::ArrayXd upperBoundConst(2 + p_cuts.cols());
        lowBoundConst(0) = p_stock(0);
        upperBoundConst(0) = p_stock(0);
        lowBoundConst(1) = p_stock(1);
        upperBoundConst(1) = p_stock(1);

        // add cuts
        for (int icut = 0; icut < p_cuts.cols(); ++icut)
        {
            double affineValue = p_cuts(0, icut);
            int ipos = 6 + 3 * icut;
            rows(ipos) =  2 + icut;
            columns(ipos) = 6;
            elements(ipos) = 1;
            // first stock
            rows(ipos + 1) = 2 + icut;
            columns(ipos + 1) = 2;
            double deriv1 =  std::max(p_cuts(1, icut), 0.);
            elements(ipos + 1) = -deriv1 ;
            // second stock
            rows(ipos + 2) = 2 + icut;
            columns(ipos + 2) = 5;
            double deriv2 =  std::max(p_cuts(2, icut), 0.);
            elements(ipos + 2) = -deriv2 ;
            lowBoundConst(2 + icut) =  -StOpt::infty;
            upperBoundConst(2 + icut) = affineValue ;
        }
        //  model
        ClpSimplex  model;
        //#ifdef NDEBUG
        model.setLogLevel(0);
        //#endif

        model.loadProblem(CoinPackedMatrix(false, rows.data(), columns.data(), elements.data(), elements.size()), lowBound.data(), upperBound.data(), objFunc.data(), lowBoundConst.data(), upperBoundConst.data());

        model.setOptimizationDirection(-1) ; // maximize

        ClpSolve solvectl;
        //solvectl.setSolveType(ClpSolve::usePrimal);
        solvectl.setSolveType(ClpSolve::useDual);
        solvectl.setPresolveType(ClpSolve::presolveOn);
        model.initialSolve(solvectl);

        bool modelSolved = model.isProvenOptimal();


        // optimal values
        p_valueAndDerivatives(0) = model.objectiveValue();

        // duals
        double *dual = model.dualRowSolution();
        p_valueAndDerivatives(1) = dual[0];
        p_valueAndDerivatives(2) = dual[1];

        // primal
        double *columnPrimal = model.primalColumnSolution();
        // for each stock
        p_stateFollowing(0) = p_stock(0) + columnPrimal[0] + columnPrimal[1];
        p_stateFollowing(1) = p_stock(1) + columnPrimal[3] + columnPrimal[4];

        // gain
        p_gain =  - (p_spot + m_injectionCost) * (columnPrimal[0] + columnPrimal[3]) - (p_spot - m_withdrawalCost) * (columnPrimal[1] + columnPrimal[4]);

        if (!modelSolved)
        {
            std::cout << "[problemLP::solveModel] : Warning Linear program could not be solved optimally somehow\n";
            abort();
        }
    }


    /// \brief defines a step in optimization
    /// \param p_stock     coordinate of the stock point to treat at current time step
    /// \param p_condEsp   continuation values  (permitting to interpolate in stocks the regresses values)
    /// \return    Gives the solution for each particle , and cut (row)
    ///            For a given simulation , cuts components (C) at a point stock \$ \bar S \f$  are given such that the cut is given by
    ///            \f$  C[0] + \sum_{i=1}^d C[i] (S_i - \bat S_i)   \f$
    Eigen::ArrayXd  stepOptimize(const Eigen::ArrayXd   &p_stock,
                                 const StOpt::ContinuationCutsGridAdaptNonConcave  &p_condEsp) const
    {
        int nbSimul = m_simulator->getNbSimul();
        // Spot price : here given by a composition of
        //  -the method getParticles from the simulator : its gives the regression factor of the models
        //  -the method fromParticlesToSpot  reconstructing the spot
        Eigen::ArrayXd spotPrice =  m_simulator->fromParticlesToSpot(m_simulator->getParticles()).array();
        // get the arrival grid
        std::shared_ptr<StOpt::GridAdaptBase> grid = p_condEsp.getGrid();
        // Attainable points : here point that can be reached are independant of the simulation
        Eigen::ArrayXXd   hypStock(2, 2);
        std::shared_ptr<StOpt::GridAdapt2D> const pGrid2D = std::dynamic_pointer_cast<StOpt::GridAdapt2D>(grid);
        hypStock(0, 0) =  std::max(p_stock(0) - m_withdrawalRate, pGrid2D->getXMin());
        hypStock(0, 1) =  std::min(p_stock(0) + m_injectionRate, pGrid2D->getXMax());
        hypStock(1, 0) =   std::max(p_stock(1) - m_withdrawalRate, pGrid2D->getYMin());
        hypStock(1, 1) =   std::min(p_stock(1) + m_injectionRate, pGrid2D->getYMax());

        // nb cuts coeff
        int nbDimCut = 3;
        // solution for return
        Eigen::ArrayXd solution(nbDimCut * nbSimul);
        // to store vales and derivatives
        Eigen::ArrayXd valueAndDerivatives(nbDimCut);
        Eigen::ArrayXd stateFollowing(p_stock.size());
        double gain ;
        // std::cout << "PTSTOCK OTIM " <<  p_stock.transpose() << "HYPOS Xmin " << hypStock(0, 0)  << " XMax " <<  hypStock(0, 1) << " YMin" << hypStock(1, 0) << " YMAx" << hypStock(1, 1) <<  std::endl ;
        // nest on  simulation
        for (int is  = 0; is < nbSimul; ++is)
        {
            // get back cuts
            std::vector< std::pair <std::shared_ptr< StOpt::GridAdaptBase>, std::shared_ptr<Eigen::ArrayXXd>  > >  meshsAndCuts = (m_bGatherConcaveGrids ? p_condEsp.getCutsConcGatherASim(hypStock, is, false) : p_condEsp.getCutsASim(hypStock, is));
            if (m_bGatherConcaveGrids)
            {
                // the case is linear so  concave to the solver should bring back one mesh
                if (meshsAndCuts.size() != 1)
                {
                    std::cout << " Gas is concave so one grid expected : got" << meshsAndCuts.size() << " So lost of concavity" <<  std::endl ;
                }
            }
            // store best result
            double bestRes = -1e10;
            // test all meshes
            for (const auto &aMeshAndCuts : meshsAndCuts)
            {

                std::shared_ptr<StOpt::GridAdapt2D> const pGrid2D = std::dynamic_pointer_cast<StOpt::GridAdapt2D>(aMeshAndCuts.first);
                double xMinLoc = std::max(pGrid2D->getXMin(), hypStock(0, 0));
                double xMaxLoc = std::min(pGrid2D->getXMax(), hypStock(0, 1));
                double yMinLoc = std::max(pGrid2D->getYMin(), hypStock(1, 0));
                double yMaxLoc = std::min(pGrid2D->getYMax(), hypStock(1, 1));
                // std::cout << "Xmin " << aMeshAndCuts.first->getXMin() << " Xmax " << aMeshAndCuts.first->getXMax() << " YMin " <<  aMeshAndCuts.first->getYMin() << "  YMax " << aMeshAndCuts.first->getYMax() << " xMINcOR " << xMinLoc << " XmAXcO" << xMaxLoc << " YmINcOR " << yMinLoc << " ymAXcOR " << yMaxLoc ;

                // Solve LP
                createAndSolveLP(*aMeshAndCuts.second, xMinLoc, xMaxLoc, yMinLoc, yMaxLoc,
                                 p_stock, spotPrice(is), valueAndDerivatives, stateFollowing, gain);
                if (valueAndDerivatives(0) > bestRes)
                {
                    bestRes =  valueAndDerivatives(0) ;
                    for (int ic = 0; ic <  nbDimCut; ++ic)
                        solution(ic * nbSimul + is, 0) = valueAndDerivatives(ic);
                }
                // std::cout << " G " << gain << " GOPT" <<   bestRes <<  " Sol" << solution.transpose() <<  std::endl ;
            }
        }
        return solution;
    }

    /// \brief get number of regimes
    inline int getNbRegime() const
    {
        return 1;
    }

    /// \brief number of controls
    inline int getNbControl() const
    {
        return 2;
    }

    /// \brief defines a step in simulation
    /// Notice that this implementation is not optimal. In fact no interpolation is necessary for this asset.
    /// This implementation is for test and example purpose
    /// \param p_continuation  defines the continuation operator for each regime
    /// \param p_state         defines the state value (modified)
    /// \param p_phiInOut      defines the value functions (modified): size number of functions to follow
    void stepSimulate(const StOpt::ContinuationCutsGridAdaptNonConcave   &p_continuation,
                      StOpt::StateWithStocks &p_state, Eigen::Ref<Eigen::ArrayXd> p_phiInOut) const
    {
        // optimal stock  attained
        Eigen::ArrayXd ptStockCur = p_state.getPtStock();
        // spot price
        double spotPrice = m_simulator->fromOneParticleToSpot(p_state.getStochasticRealization());

        // get the arrival grid
        std::shared_ptr<StOpt::GridAdaptBase> grid = p_continuation.getGrid();

        // Attainable points : here point that can be reached are independant of the simulation
        Eigen::ArrayXXd   hypStock(2, 2);
        std::shared_ptr<StOpt::GridAdapt2D> const pGrid2D = std::dynamic_pointer_cast<StOpt::GridAdapt2D>(grid);
        hypStock(0, 0) =  std::max(ptStockCur(0) - m_withdrawalRate, pGrid2D->getXMin());
        hypStock(0, 1) =  std::min(ptStockCur(0) + m_injectionRate, pGrid2D->getXMax());
        hypStock(1, 0) =   std::max(ptStockCur(1) - m_withdrawalRate, pGrid2D->getYMin());
        hypStock(1, 1) =   std::min(ptStockCur(1) + m_injectionRate, pGrid2D->getYMax());
        // mesh and cuts
        std::vector< std::pair <std::shared_ptr< StOpt::GridAdaptBase>, std::shared_ptr<Eigen::ArrayXXd>  > >  meshsAndCuts = (m_bGatherConcaveGrids ? p_continuation.getCutsConcGatherASim(hypStock, p_state.getStochasticRealization(), false) :
                p_continuation.getCutsASim(hypStock, p_state.getStochasticRealization()));

        // nb cuts
        int nbDimCut = 3;
        // to store vales and derivatives
        Eigen::ArrayXd valueAndDerivatives(nbDimCut);
        Eigen::ArrayXd stateFollowing(2);
        double gainOpt ;
        // store best result
        double bestRes = -1e10;
        // test all meshes
        for (const auto &aMeshAndCuts : meshsAndCuts)
        {
            double gain ;

            std::shared_ptr<StOpt::GridAdapt2D> const pGrid2D = std::dynamic_pointer_cast<StOpt::GridAdapt2D>(aMeshAndCuts.first);
            double xMinLoc = std::max(pGrid2D->getXMin(), hypStock(0, 0));
            double xMaxLoc = std::min(pGrid2D->getXMax(), hypStock(0, 1));
            double yMinLoc = std::max(pGrid2D->getYMin(), hypStock(1, 0));
            double yMaxLoc = std::min(pGrid2D->getYMax(), hypStock(1, 1));

            // Solve LP
            createAndSolveLP(*aMeshAndCuts.second,  xMinLoc, xMaxLoc, yMinLoc, yMaxLoc,
                             ptStockCur, spotPrice, valueAndDerivatives, stateFollowing, gain);

            if (valueAndDerivatives(0) > bestRes)
            {
                bestRes = valueAndDerivatives(0);
                gainOpt = gain;
                p_state.setPtStock(stateFollowing);
            }
        }
        p_phiInOut(0) += gainOpt ;
    }

    // here  refine the grids
    bool refineGrid(std::shared_ptr<StOpt::GridAdaptBase>   &p_gridToAdapt, const  Eigen::ArrayXXd &pCuts)
    {
        if (!m_bRefine)
            return false;

        int istep = static_cast<int>(m_simulator->getStep());
        if (istep != m_iStep)
        {
            m_iStep = istep;
            // test just refine once all the meshes
            std::shared_ptr<StOpt::GridAdapt2D> const pGrid2D = std::dynamic_pointer_cast<StOpt::GridAdapt2D>(p_gridToAdapt);
            std::list< std::pair<std::shared_ptr< StOpt::Mesh2D>, std::shared_ptr< std::vector< Eigen::ArrayXi > > > >   meshesAndLevel =   pGrid2D->getMeshes();
            for (auto &meshALevel : meshesAndLevel)
                pGrid2D->splitMesh(meshALevel);

            return true;
        }
        return false;
    }
    // void setRefinementToFalse() {m_bRaff = false;}

    ///\brief store the simulator
    inline void setSimulator(const std::shared_ptr<Simulator> &p_simulator)
    {
        m_simulator = p_simulator ;
    }

    /// \brief get the simulator back
    inline std::shared_ptr< StOpt::SimulatorDPBase > getSimulator() const
    {
        return m_simulator ;
    }

    /// \brief get size of the  function to follow in simulation
    inline int getSimuFuncSize() const
    {
        return 1;
    }
}
;
#endif /* OPTIMIZEGASSTORAGECUT_H */
