// Copyright (C) 2021 EDF
// All Rights Reserved
// This code is published under the GNU Lesser General Public License (GNU LGPL)
#ifndef OPTIMIZESTORAGE_H
#define OPTIMIZESTORAGE_H
#include <memory>
#include <Eigen/Dense>
#include <StOpt/dp/OptimizerDPBase.h>
#include <StOpt/core/utils/StateWithStocks.h>
#include <StOpt/regression/BaseRegression.h>
#include <StOpt/core/grids/Interpolator.h>
#include <StOpt/core/utils/constant.h>
#include <StOpt/regression/ContinuationValue.h>
#include "StOpt/regression/GridAndRegressedValue.h"
#include "StOpt/regression/GridAndRegressedValueAutoDiff.h"
#include "test/c++/tools/utils/Capa.h"


/** \file OptimizeStorage.h
 * Optimize storage commands
 * Injection and withdrawal are given by time step
 * \author Xavier Warin
 */
template<class Simulator>
class OptimizeStorage : public StOpt::OptimizerDPBase
{
private:

    const Capa   &m_injectionVolume ; ///< Injection volume  capacity
    const Capa   &m_withDrawalVolume ; ///< Withdrawal volume  capacity
    const double m_alphaADiff; // autodiff regularization

    /// store the simulator
    std::shared_ptr<Simulator> m_simulator;


public:

    /// \brief Constructor
    /// \param p_injectionVolume   Injection volume capacity
    /// \param p_withDrawalVolume  Withdrawal volume  capacity
    OptimizeStorage(const Capa   &p_injectionVolume, const Capa &p_withDrawalVolume, double p_alphaADiff = 20.): m_injectionVolume(p_injectionVolume), m_withDrawalVolume(p_withDrawalVolume), m_alphaADiff(p_alphaADiff) {}

    /// \brief define the diffusion cone for parallelism (not needed if not parallelism required )
    /// \return returns in each dimension the min max values in the stock that can be reached from the grid p_gridByProcessor for each regime
    std::vector< std::array< double, 2> > getCone(const  std::vector<  std::array< double, 2>  > &) const
    {
        std::vector< std::array< double, 2> > extrGrid(1);
        extrGrid[0][0] = - StOpt::infty;
        extrGrid[0][1] = StOpt::infty;
        return extrGrid;
    }

    /// \brief defines the dimension to split for MPI parallelism (not neeede if no parallelism )
    ///        For each dimension return true is the direction can be split
    Eigen::Array< bool, Eigen::Dynamic, 1> getDimensionToSplit() const
    {
        Eigen::Array< bool, Eigen::Dynamic, 1> bDim = Eigen::Array< bool, Eigen::Dynamic, 1>::Constant(1, true);
        return  bDim ;
    }

    /// \brief defines a step in optimization
    /// \param p_grid      grid at arrival step after command
    /// \param p_stock     coordinate of the stock point to treat at current time step
    /// \param p_condEsp   continuation values for each regime (permitting to interpolate in stocks the regresses values)
    /// \param p_phiIn     for each regime  gives the solution calculated at the previous step ( next time step by Dynamic Programming resolution)
    /// \return   a pair  :
    ///              - for each regimes (column) gives the solution for each particle (row)
    ///              - for each control (column) gives the optimal control for each particle (rows)
    ///              .
    std::pair< Eigen::ArrayXXd, Eigen::ArrayXXd> stepOptimize(const   std::shared_ptr< StOpt::SpaceGrid> &p_grid, const Eigen::ArrayXd   &p_stock,
            const std::vector<StOpt::ContinuationValue> &p_condEsp,
            const std::vector < std::shared_ptr< Eigen::ArrayXXd > > &p_phiIn) const
    {
        // std::shared_ptr<MeanRevertingNDSimulator> ptSimul=  std::static_pointer_cast<MeanRevertingNDSimulator>( m_simulator);
        int nbSimul = m_simulator->getNbSimul();
        std::pair< Eigen::ArrayXXd, Eigen::ArrayXXd> solutionAndControl;
        solutionAndControl.first.resize(nbSimul, 1);
        solutionAndControl.second.resize(nbSimul, 1);
        // get back spot
        Eigen::ArrayXd spot = m_simulator->getSpot();
        // get current time Step number : here simulator
        int iStep  = static_cast<int>(m_simulator->getCurrentStep() + StOpt::tiny);
        //  injection Max
        double maxStorage = p_grid->getExtremeValues()[0][1];
        double injectionRate = m_injectionVolume.getValue(iStep, p_stock(0));
        double injectionMax = std::min(maxStorage - p_stock(0) - StOpt::tiny, injectionRate);
        // withdrawl max (algebric so negative)
        double minStorage = p_grid->getExtremeValues()[0][0];
        double withdrawalRate = m_withDrawalVolume.getValue(iStep, p_stock(0));
        double withdrawalMax = std::min(p_stock(0) - minStorage - StOpt::tiny, withdrawalRate);

        // 2 comparisons to avoid withdrawalMax= withdrawalMax =0
        double smallDec = StOpt::tiny * std::max(std::max(std::fabs(withdrawalMax), std::fabs(injectionMax)), 2.);
        // test if point is possible
        if (-withdrawalMax > (injectionMax + smallDec))
        {
            solutionAndControl.first.col(0)  = Eigen::ArrayXd::Constant(nbSimul, - StOpt::infty);
            solutionAndControl.second.col(0) = Eigen::ArrayXd::Constant(nbSimul, - StOpt::infty);
            std::cout << " IMPOS iStep " << iStep <<  "Pt " << p_stock(0) << " WITH " << withdrawalMax << " injectionMax  " << injectionMax <<  " diff " << withdrawalMax + injectionMax << " minStorage Next step " << minStorage << " max storage Next step " <<  maxStorage << " injectionRate " <<  injectionRate << " withdrawalRate " << withdrawalRate << std::endl ;
            return solutionAndControl;
        }

        // cash, optimal command, condi
        Eigen::ArrayXd cashOpt = Eigen::ArrayXd::Constant(nbSimul, - StOpt::infty);
        Eigen::ArrayXd commandOpt = Eigen::ArrayXd::Constant(nbSimul, - StOpt::infty);
        Eigen::ArrayXd condOpt = Eigen::ArrayXd::Constant(nbSimul, - StOpt::infty);
        // test max withdrawal command
        {
            double command = -withdrawalMax;
            Eigen::ArrayXd nextStock = p_stock + command;
            // interpolator next pt
            std::shared_ptr<StOpt::Interpolator<double>>  interpolatorNextStock = p_grid->createInterpolator(nextStock);
            // interpolate stock
            Eigen::ArrayXd cashNext = interpolatorNextStock->applyVec(*p_phiIn[0]);
            // conditional expectation
            Eigen::ArrayXd condExp = - command * spot +  p_condEsp[0].getAllSimulations(*interpolatorNextStock);
            for (int is = 0; is < nbSimul; ++is)
            {
                if (condExp(is) > condOpt(is))
                {
                    condOpt(is) = condExp(is);
                    cashOpt(is) = - command *  spot(is) + cashNext(is);
                    commandOpt(is) = command;
                }
            }
        }
        // test max injection  command
        {
            double command = injectionMax;
            Eigen::ArrayXd nextStock = p_stock + command;
            // interpolator next pt
            std::shared_ptr<StOpt::Interpolator<double>>  interpolatorNextStock = p_grid->createInterpolator(nextStock);
            // interpolate stock
            Eigen::ArrayXd cashNext = interpolatorNextStock->applyVec(*p_phiIn[0]);
            // conditional expectation
            Eigen::ArrayXd condExp = - command * spot +  p_condEsp[0].getAllSimulations(*interpolatorNextStock);
            for (int is = 0; is < nbSimul; ++is)
            {
                if (condExp(is) > condOpt(is))
                {
                    condOpt(is) = condExp(is);
                    cashOpt(is) = - command *  spot(is) + cashNext(is);
                    commandOpt(is) = command;
                }
            }
        }
        double stockMaxReached = p_stock(0) + injectionMax;
        double stockMinReached = p_stock(0) - withdrawalMax;
        // test interor points
        Eigen::ArrayXd grid1D  = *(std::static_pointer_cast<StOpt::FullGrid>(p_grid)->get1DGrids()[0]);
        int iMin = 0;
        while (grid1D(iMin) <= stockMinReached)
        {
            iMin += 1;
            if (iMin == (grid1D.size() - 1))
                break;
        }
        int iMax = grid1D.size() - 1;
        while (grid1D(iMax) >= stockMaxReached)
        {
            iMax -= 1;
            if (iMax == 0)
                break;
        }
        if (iMin <= iMax)
        {
            for (int ipt = iMin; ipt <= iMax; ++ipt)
            {
                if (std::fabs(grid1D(ipt)) > StOpt::tiny)
                {
                    // command
                    double command = grid1D(ipt) - p_stock(0);
                    Eigen::ArrayXd nextStock = p_stock + command;
                    // interpolator next pt
                    std::shared_ptr<StOpt::Interpolator<double>>  interpolatorNextStock = p_grid->createInterpolator(nextStock);
                    // interpolate stock
                    Eigen::ArrayXd cashNext = interpolatorNextStock->applyVec(*p_phiIn[0]);
                    // conditional expectation
                    Eigen::ArrayXd condExp = - command * spot +  p_condEsp[0].getAllSimulations(*interpolatorNextStock);
                    for (int is = 0; is < nbSimul; ++is)
                    {
                        if (condExp(is) > condOpt(is))
                        {
                            condOpt(is) = condExp(is);
                            cashOpt(is) = - command *  spot(is) + cashNext(is);
                            commandOpt(is) = command;
                        }
                    }
                }
            }
        }
        for (int is = 0; is < nbSimul; ++is)
        {
            solutionAndControl.first(is, 0)  = cashOpt(is);
            solutionAndControl.second(is, 0) = commandOpt(is);
        }
        return solutionAndControl;
    }

    /// \brief defines a step in simulation
    /// This implementation is for test and example purpose
    /// \param p_grid          grid at arrival step after command
    /// \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 std::shared_ptr< StOpt::SpaceGrid>   &p_grid, const std::vector< StOpt::GridAndRegressedValue  > &p_continuation,
                      StOpt::StateWithStocks<double> &p_state, Eigen::Ref<Eigen::ArrayXd> p_phiInOut) const
    {
        //std::shared_ptr<MeanRevertingNDSimulator> ptSimul=  std::static_pointer_cast<MeanRevertingNDSimulator>( m_simulator);
        // optimal stock  attained
        Eigen::ArrayXd ptStockCur = p_state.getPtStock();
        Eigen::ArrayXd ptStockMax(ptStockCur);
        // spot price
        double spot = m_simulator->getSpotUncertainties(p_state.getStochasticRealization());
        // std::cout << " SimReop spot " << spot << "  OU " << p_state.getStochasticRealization().transpose() << std::endl ;
        // size of the stock
        double minStorage = p_grid->getExtremeValues()[0][0];
        double maxStorage = p_grid->getExtremeValues()[0][1];
        // get current time Step number : here simulator
        int iStep  = static_cast<int>(m_simulator->getCurrentStep() + StOpt::tiny);
        // if injection
        double injectionRate = m_injectionVolume.getValue(iStep, ptStockCur(0));
        double injectionMax = std::min(maxStorage - p_state.getPtStock()(0) - StOpt::tiny, injectionRate);
        double withdrawalRate = m_withDrawalVolume.getValue(iStep, ptStockCur(0));
        double withdrawalMax = std::min(p_state.getPtStock()(0) - minStorage - StOpt::tiny, withdrawalRate);
        double  cashOpt = -1e30;
        double  commandOpt =  -1e30;
        double  condOpt =  -1e30;
        // test max withdrawal
        {
            double command = -withdrawalMax;
            Eigen::ArrayXd nextStock = ptStockCur + command;
            double gainLoc = - command *  spot;
            double continuation =  p_continuation[0].getValue(nextStock, p_state.getStochasticRealization());
            double espCond = gainLoc + continuation ;
            if (espCond > condOpt)
            {
                condOpt = espCond;
                cashOpt = gainLoc ;
                commandOpt = command;
            }
        }
        {
            double command = injectionMax;
            Eigen::ArrayXd nextStock = ptStockCur + command;
            double gainLoc = - command *  spot;
            double continuation =  p_continuation[0].getValue(nextStock, p_state.getStochasticRealization());
            double espCond = gainLoc + continuation ;
            if (espCond > condOpt)
            {
                condOpt = espCond;
                cashOpt = gainLoc ;
                commandOpt = command;
            }
        }
        double stockMaxReached = ptStockCur(0) + injectionMax;
        double stockMinReached = ptStockCur(0) - withdrawalMax;
        Eigen::ArrayXd grid1D  = *(std::static_pointer_cast<StOpt::FullGrid>(p_grid)->get1DGrids()[0]);

        int iMin = 0;
        while (grid1D(iMin) <= stockMinReached)
        {
            iMin += 1;
            if (iMin == (grid1D.size() - 1))
                break;
        }
        int iMax = grid1D.size() - 1;
        while (grid1D(iMax) >= stockMaxReached)
        {
            iMax -= 1;
            if (iMax == 0)
                break;
        }
        if (iMin <= iMax)
        {
            for (int ipt = iMin; ipt <= iMax; ++ipt)
            {
                if (std::fabs(grid1D(ipt)) > StOpt::tiny)
                {
                    // command
                    double command = grid1D(ipt) - ptStockCur(0);
                    Eigen::ArrayXd nextStock = ptStockCur + command;
                    double gainLoc = - command *  spot;
                    double continuation =  p_continuation[0].getValue(nextStock, p_state.getStochasticRealization());
                    double espCond = gainLoc + continuation ;
                    if (espCond > condOpt)
                    {
                        condOpt = espCond;
                        cashOpt = gainLoc ;
                        commandOpt = command;
                    }
                }
            }
        }
        // for return
        p_state.setPtStock(ptStockCur + commandOpt);
        p_phiInOut(0) += cashOpt ;
        p_phiInOut(1) = commandOpt ;
        p_phiInOut(2) = -commandOpt * spot;
    }

    /// \brief Defines a step in simulation using interpolation in controls
    /// \param p_grid          grid at arrival step after command
    /// \param p_control       defines the controls
    /// \param p_state         defines the state value (modified)
    /// \param p_phiInOut      defines the value function (modified): size number of functions to follow
    virtual void stepSimulateControl(const std::shared_ptr< StOpt::SpaceGrid>   &p_grid, const std::vector< StOpt::GridAndRegressedValue  > &p_control,
                                     StOpt::StateWithStocks<double> &p_state,
                                     Eigen::Ref<Eigen::ArrayXd>  p_phiInOut) const
    {
        // cast
        //std::shared_ptr<MeanRevertingNDSimulator> ptSimul=  std::static_pointer_cast<MeanRevertingNDSimulator>( m_simulator);
        Eigen::ArrayXd ptStock = p_state.getPtStock();
        // spot price
        double spotPrice = m_simulator->getSpotUncertainties(p_state.getStochasticRealization());
        // std::cout << " SimCont spot " << spotPrice << "  OU " << p_state.getStochasticRealization().transpose() << std::endl ;
        // optimal control
        double control = p_control[0].getValue(p_state.getPtStock(), p_state.getStochasticRealization());
        double maxStorage = p_grid->getExtremeValues()[0][1];
        double minStorage =  p_grid->getExtremeValues()[0][0];
        control = std::max(std::min(maxStorage - ptStock(0), control), minStorage - ptStock(0));
        p_phiInOut(0) -= control * spotPrice;
        ptStock(0) += control ;
        ptStock(0) = std::min(std::max(minStorage, ptStock(0)), maxStorage); // avoid rounding error
        p_state.setPtStock(ptStock);
        p_phiInOut(1) = control ;
        p_phiInOut(2) = -control * spotPrice;
        double tangent =  m_simulator->getTangent(p_state.getStochasticRealization());
        p_phiInOut(3) = -control * tangent;
        p_phiInOut(4) = -control * spotPrice * (m_simulator->getTrendDeriv(0) + p_state.getStochasticRealization()(0));
        p_phiInOut(5) = -control * spotPrice * (m_simulator->getTrendDeriv(1) + p_state.getStochasticRealization()(1));
    }



    template<typename T, typename Expr>
    void stepSimulateControlAutoDiff(const std::shared_ptr< StOpt::FullGrid>   &p_grid, const std::vector< StOpt::GridAndRegressedValueAutoDiff<T>  > &p_control,
                                     StOpt::StateWithStocks<T> &p_state,
                                     Expr  &&p_phiInOut)
    {
        // foward reference to temporary
        auto &&phiInOut = p_phiInOut;
        // // cast
        // std::shared_ptr<MeanRevertingNDSimulatorT<T>> ptSimul=  std::static_pointer_cast<MeanRevertingNDSimulatorT<T>>( m_simulator);
        // Templated code using the autodiff type
        Eigen::Array<T, Eigen::Dynamic, 1> ptStock = p_state.getPtStock();
        // spot price
        T spotPrice = m_simulator->getSpotUncertaintiesAD(p_state.getStochasticRealization());
        // std::cout << " SimContAD spot " << spotPrice << "  OU " << p_state.getStochasticRealization().transpose() << std::endl ;
        // optimal control
        T control = p_control[0].getValueAutoDiff(p_state.getPtStock(), p_state.getStochasticRealization());
        T maxStorage = T(p_grid->getExtremeValues()[0][1]);
        T minStorage =  T(p_grid->getExtremeValues()[0][0]);
        T upperControl = maxStorage - ptStock(0);
        T lowerControl = minStorage - ptStock(0);
        using StOpt::toDouble;
        if (toDouble(control) > toDouble(upperControl))
            control = upperControl;
        if (toDouble(control) < toDouble(lowerControl))
            control = lowerControl;
        phiInOut(0) -= control * spotPrice;
        ptStock(0) += control ;
        p_state.setPtStock(ptStock);
        p_phiInOut(1) = control ;
        p_phiInOut(2) = -control * spotPrice;
    }



    ///\brief store the simulator
    inline void setSimulator(const std::shared_ptr<StOpt::SimulatorDPBase> &p_simulator)
    {
        m_simulator = std::static_pointer_cast<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 6;
    }

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

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