// Copyright (C) 2019 EDF
// All Rights Reserved
// This code is published under the GNU Lesser General Public License (GNU LGPL)

/** \file Pybind11StOptGlobal.cpp
 * \brief Mapping optimizer and simulator  classes to python
 * \author Xavier Warin
 */
#include <Eigen/Dense>
#include <memory>
#include <functional>
#include <array>
#ifdef USE_MPI
#include <boost/mpi.hpp>
#endif
#include <pybind11/pybind11.h>
#include <pybind11/eigen.h>
#include <pybind11/stl_bind.h>
#include <pybind11/stl.h>
#include "StOpt/dp/SimulatorDPBase.h"
#include "StOpt/dp/SimulatorDPBaseTree.h"
#include "StOpt/dp/OptimizerDPBase.h"
#include "StOpt/dp/OptimizerDPTreeBase.h"
#include "StOpt/dp/OptimizerSwitchBase.h"
#include "StOpt/core/grids/SpaceGrid.h"
#include "StOpt/core/grids/FullGrid.h"
#include "StOpt/core/grids/RegularSpaceIntGrid.h"
#include "StOpt/core/utils/StateWithStocks.h"
#include "StOpt/core/utils/StateWithIntState.h"
#include "StOpt/tree/StateTreeStocks.h"
#include "StOpt/tree/ContinuationValueTree.h"
#include "StOpt/regression/BaseRegression.h"
#include "StOpt/dp/FinalStepDP.h"
#include "StOpt/dp/TransitionStepRegressionDP.h"
#include "StOpt/dp/TransitionStepRegressionSwitch.h"
#include "StOpt/dp/SimulateStepRegression.h"
#include "StOpt/dp/SimulateStepRegressionControl.h"
#include "StOpt/dp/SimulateStepSwitch.h"
#include "StOpt/dp/TransitionStepTreeDP.h"
#include "StOpt/dp/SimulateStepTree.h"
#include "StOpt/dp/SimulateStepTreeControl.h"
#include "StOpt/semilagrangien/SemiLagrangEspCond.h"
#ifdef USE_MPI
#include "StOpt/dp/FinalStepDPDist.h"
#include "StOpt/dp/FinalStepZeroDist.h"
#include "StOpt/dp/TransitionStepRegressionDPDist.h"
#include "StOpt/dp/TransitionStepRegressionSwitchDist.h"
#include "StOpt/dp/SimulateStepRegressionDist.h"
#include "StOpt/dp/SimulateStepRegressionControlDist.h"
#include "StOpt/dp/TransitionStepTreeDPDist.h"
#include "StOpt/dp/SimulateStepTreeDist.h"
#include "StOpt/dp/SimulateStepTreeControlDist.h"
#include "StOpt/core/parallelism/reconstructProc0Mpi.h"
#include "StOpt/core/parallelism/reconstructProc0ForIntMpi.h"
#endif
#include "StOpt/python/BinaryFileArchiveStOpt.h"
#include "StOpt/python/PayOffBase.h"
#include  "StOpt/python/Pybind11VectorAndList.h"

namespace py = pybind11;


//Wrapper for simulator of regression
//***********************************

class PySimulatorDPBase : public StOpt::SimulatorDPBase
{
public:

    using StOpt::SimulatorDPBase::SimulatorDPBase;

    Eigen::MatrixXd getParticles() const override
    {
        PYBIND11_OVERLOAD_PURE(Eigen::MatrixXd, StOpt::SimulatorDPBase, getParticles);
    }

    void stepForward() override
    {
        PYBIND11_OVERLOAD_PURE(void,  StOpt::SimulatorDPBase, stepForward) ;
    }
    void stepBackward() override
    {
        PYBIND11_OVERLOAD_PURE(void,  StOpt::SimulatorDPBase, stepBackward);
    }

    Eigen::MatrixXd  stepForwardAndGetParticles() override
    {
        PYBIND11_OVERLOAD_PURE(Eigen::MatrixXd,  StOpt::SimulatorDPBase, stepForwardAndGetParticles);
    }

    Eigen::MatrixXd  stepBackwardAndGetParticles() override
    {
        PYBIND11_OVERLOAD_PURE(Eigen::MatrixXd,  StOpt::SimulatorDPBase, stepBackwardAndGetParticles);
    }

    int getDimension() const override
    {
        PYBIND11_OVERLOAD_PURE(int, StOpt::SimulatorDPBase, getDimension);
    }
    int getNbStep() const  override
    {
        PYBIND11_OVERLOAD_PURE(int, StOpt::SimulatorDPBase, getNbStep);
    }
    double getStep() const  override
    {
        PYBIND11_OVERLOAD_PURE(double, StOpt::SimulatorDPBase, getStep) ;
    }
    int getNbSimul() const  override
    {
        PYBIND11_OVERLOAD_PURE(int, StOpt::SimulatorDPBase, getNbSimul);
    }
    double getCurrentStep() const  override
    {
        PYBIND11_OVERLOAD_PURE(double, StOpt::SimulatorDPBase, getCurrentStep);
    }
    double  getActuStep() const  override
    {
        PYBIND11_OVERLOAD_PURE(double, StOpt::SimulatorDPBase, getActuStep);
    }
    double  getActu() const override
    {
        PYBIND11_OVERLOAD_PURE(double, StOpt::SimulatorDPBase, getActu);
    }
};

//Wrapper for simulator of tree
//***********************************

class PySimulatorDPBaseTree : public StOpt::SimulatorDPBaseTree
{
public:

    using StOpt::SimulatorDPBaseTree::SimulatorDPBaseTree;

    void stepForward() override
    {
        PYBIND11_OVERLOAD_PURE(void,  StOpt::SimulatorDPBaseTree, stepForward) ;
    }
    void stepBackward() override
    {
        PYBIND11_OVERLOAD_PURE(void,  StOpt::SimulatorDPBaseTree, stepBackward);
    }
    int getNbSimul() const  override
    {
        PYBIND11_OVERLOAD_PURE(int, StOpt::SimulatorDPBaseTree, getNbSimul);
    }
    Eigen::ArrayXd  getValueAssociatedToNode(const int &) const override
    {
        PYBIND11_OVERLOAD_PURE(Eigen::ArrayXd, StOpt::SimulatorDPBaseTree, getValueAssociatedToNode);
    }

    int getNodeAssociatedToSim(const int &) const override
    {
        PYBIND11_OVERLOAD_PURE(int, StOpt::SimulatorDPBaseTree, getNodeAssociatedToSim);
    }
};


///wrapper for optimization object for Optimization Base object for tree for dynamic programming
///**************************************************************************************************
class  PyOptimizerDPTreeBase :  public StOpt::OptimizerDPTreeBase
{
    using StOpt::OptimizerDPTreeBase::OptimizerDPTreeBase;

    std::vector< std::array< double, 2> > getCone(const  std::vector<  std::array< double, 2>  > &p_regionByProcessor) const override
    {
        using vecTep = std::vector< std::array< double, 2> >;
        PYBIND11_OVERLOAD_PURE(vecTep, StOpt::OptimizerDPTreeBase, getCone, p_regionByProcessor);
    }

    Eigen::Array< bool, Eigen::Dynamic, 1> getDimensionToSplit() const override
    {
        using boolArray = Eigen::Array< bool, Eigen::Dynamic, 1> ;
        PYBIND11_OVERLOAD_PURE(boolArray, StOpt::OptimizerDPTreeBase, getDimensionToSplit);
    }

    std::pair< Eigen::ArrayXXd, Eigen::ArrayXXd>   stepOptimize(const std::shared_ptr< StOpt::SpaceGrid> &,
            const Eigen::ArrayXd &, const std::vector<StOpt::ContinuationValueTree> &) const override
    {
        using pairArray = std::pair< Eigen::ArrayXXd, Eigen::ArrayXXd>;
        PYBIND11_OVERLOAD_PURE(pairArray, StOpt::OptimizerDPTreeBase, stepOptimize);
    }

    void stepSimulate(const std::shared_ptr< StOpt::SpaceGrid>   &p_grid, const std::vector< StOpt::GridTreeValue  > &p_continuation,
                      StOpt::StateTreeStocks &p_state,
                      Eigen::Ref<Eigen::ArrayXd> p_phiInOut) const override
    {
        PYBIND11_OVERLOAD_PURE(void, StOpt::OptimizerDPTreeBase, stepSimulate, p_grid, p_continuation, p_state, p_phiInOut);
    }

    void stepSimulateControl(const std::shared_ptr< StOpt::SpaceGrid>   &p_grid, const std::vector< StOpt::GridTreeValue  > &p_control,
                             StOpt::StateTreeStocks &p_state,
                             Eigen::Ref<Eigen::ArrayXd> p_phiInOut) const override
    {
        PYBIND11_OVERLOAD_PURE(void, StOpt::OptimizerDPTreeBase, stepSimulateControl, p_grid, p_control, p_state, p_phiInOut);
    }

    int getNbRegime() const override
    {
        PYBIND11_OVERLOAD_PURE(int,  StOpt::OptimizerDPTreeBase, getNbRegime);
    }

    int getNbControl() const override
    {
        PYBIND11_OVERLOAD_PURE(int,  StOpt::OptimizerDPTreeBase, getNbControl);
    }

    std::shared_ptr< StOpt::SimulatorDPBaseTree > getSimulator() const override
    {
        PYBIND11_OVERLOAD_PURE(std::shared_ptr< StOpt::SimulatorDPBaseTree >,  StOpt::OptimizerDPTreeBase, getSimulator);
    }

    int getSimuFuncSize() const  override
    {
        PYBIND11_OVERLOAD_PURE(int, StOpt::OptimizerDPTreeBase, getSimuFuncSize);
    }
};

///wrapper for optimization object for Optimization Base object for regression for dynamic programming
///**************************************************************************************************
class  PyOptimizerDPBase :  public StOpt::OptimizerDPBase
{
    using StOpt::OptimizerDPBase::OptimizerDPBase;

    std::vector< std::array< double, 2> > getCone(const  std::vector<  std::array< double, 2>  > &p_regionByProcessor) const override
    {
        using vecTep = std::vector< std::array< double, 2> >;
        PYBIND11_OVERLOAD_PURE(vecTep, StOpt::OptimizerDPBase, getCone, p_regionByProcessor);
    }

    Eigen::Array< bool, Eigen::Dynamic, 1> getDimensionToSplit() const override
    {
        using boolArray = Eigen::Array< bool, Eigen::Dynamic, 1> ;
        PYBIND11_OVERLOAD_PURE(boolArray, StOpt::OptimizerDPBase, getDimensionToSplit);
    }

    std::pair< Eigen::ArrayXXd, Eigen::ArrayXXd>   stepOptimize(const std::shared_ptr< StOpt::SpaceGrid> &,
            const Eigen::ArrayXd &, const std::vector<StOpt::ContinuationValue> &,
            const std::vector < std::shared_ptr< Eigen::ArrayXXd > > &) const override
    {
        using pairArray = std::pair< Eigen::ArrayXXd, Eigen::ArrayXXd>;
        PYBIND11_OVERLOAD_PURE(pairArray, StOpt::OptimizerDPBase, stepOptimize);
    }

    void stepSimulate(const std::shared_ptr< StOpt::SpaceGrid>   &p_grid, const std::vector< StOpt::GridAndRegressedValue  > &p_continuation,
                      StOpt::StateWithStocks &p_state,
                      Eigen::Ref<Eigen::ArrayXd> p_phiInOut) const override
    {
        PYBIND11_OVERLOAD_PURE(void, StOpt::OptimizerDPBase, stepSimulate, p_grid, p_continuation, p_state, p_phiInOut);
    }

    void stepSimulateControl(const std::shared_ptr< StOpt::SpaceGrid>   &p_grid, const std::vector< StOpt::GridAndRegressedValue  > &p_control,
                             StOpt::StateWithStocks &p_state,
                             Eigen::Ref<Eigen::ArrayXd> p_phiInOut) const override
    {
        PYBIND11_OVERLOAD_PURE(void, StOpt::OptimizerDPBase, stepSimulateControl, p_grid, p_control, p_state, p_phiInOut);
    }

    int getNbRegime() const override
    {
        PYBIND11_OVERLOAD_PURE(int,  StOpt::OptimizerDPBase, getNbRegime);
    }

    int getNbControl() const override
    {
        PYBIND11_OVERLOAD_PURE(int,  StOpt::OptimizerDPBase, getNbControl);
    }

    std::shared_ptr< StOpt::SimulatorDPBase > getSimulator() const override
    {
        //using sharedPtrSim=std::shared_ptr< StOpt::SimulatorDPBase >
        PYBIND11_OVERLOAD_PURE(std::shared_ptr< StOpt::SimulatorDPBase >,  StOpt::OptimizerDPBase, getSimulator);
    }

    int getSimuFuncSize() const  override
    {
        PYBIND11_OVERLOAD_PURE(int, StOpt::OptimizerDPBase, getSimuFuncSize);
    }
};


///wrapper for optimization object for Optimization Switching Base object for regression for dynamic programming
///**************************************************************************************************
class  PyOptimizerSwitchBase :  public StOpt::OptimizerSwitchBase
{
    using StOpt::OptimizerSwitchBase::OptimizerSwitchBase;

    std::vector< std::array<int, 2 > > getCone(const int &, const   std::vector< std::array<int, 2 > > &) const override
    {
        using retCone = std::vector< std::array<int, 2 > > ;
        PYBIND11_OVERLOAD_PURE(retCone, StOpt::OptimizerSwitchBase, getCone);
    }

    std::vector< Eigen::Array< bool, Eigen::Dynamic, 1> >  getDimensionToSplit() const override
    {
        using vecBoolArray = std::vector< Eigen::Array< bool, Eigen::Dynamic, 1> > ;
        PYBIND11_OVERLOAD_PURE(vecBoolArray, StOpt::OptimizerSwitchBase, getDimensionToSplit);
    }

    Eigen::ArrayXd    stepOptimize(const   std::vector<std::shared_ptr< StOpt::RegularSpaceIntGrid> > &,
                                   const   int &,
                                   const   Eigen::ArrayXi &,
                                   const   std::shared_ptr< StOpt::BaseRegression> &,
                                   const   std::vector < std::shared_ptr< Eigen::ArrayXXd > > &) const override
    {
        PYBIND11_OVERLOAD_PURE(Eigen::ArrayXd, StOpt::OptimizerSwitchBase, stepOptimize);
    }

    void stepSimulate(const std::vector< std::shared_ptr< StOpt::RegularSpaceIntGrid> > &,
                      const std::shared_ptr< StOpt::BaseRegression> &,
                      const std::vector< Eigen::ArrayXXd > &,
                      StOpt::StateWithIntState &,
                      Eigen::Ref<Eigen::ArrayXd>) const override
    {
        PYBIND11_OVERLOAD_PURE(void, StOpt::OptimizerSwitchBase, stepSimulate) ;//, p_grid, p_condExp, p_basisFunc p_state, p_phiInOut);
    }


    int getNbRegime() const override
    {
        PYBIND11_OVERLOAD_PURE(int,  StOpt::OptimizerSwitchBase, getNbRegime);
    }


    std::shared_ptr< StOpt::SimulatorDPBase > getSimulator() const override
    {
        //using sharedPtrSim=std::shared_ptr< StOpt::SimulatorSwitchBase >
        PYBIND11_OVERLOAD_PURE(std::shared_ptr< StOpt::SimulatorDPBase >,  StOpt::OptimizerSwitchBase, getSimulator);
    }

    int getSimuFuncSize() const  override
    {
        PYBIND11_OVERLOAD_PURE(int, StOpt::OptimizerSwitchBase, getSimuFuncSize);
    }
};


// wrapper for Final Step
//***********************
class PyFinalStepDP: StOpt::FinalStepDP
{
public:

    using StOpt::FinalStepDP::FinalStepDP;

    std::vector< Eigen::ArrayXXd >   operator()(py::object &p_funcValue, const Eigen::ArrayXXd &p_particles) const
    {
        auto lambda = [p_funcValue](const int &p_i, const Eigen::ArrayXd & p_x, const Eigen::ArrayXd & p_y)->double { return p_funcValue(p_i, p_x, p_y).cast<double>();};

        std::vector< std::shared_ptr<Eigen::ArrayXXd>  >  resLoc = StOpt::FinalStepDP::operator()(std::function<double(const int &, const Eigen::ArrayXd &, const Eigen::ArrayXd &)> (lambda), p_particles);
        std::vector< Eigen::ArrayXXd > res;
        res.reserve(resLoc.size());
        for (const auto   &item : resLoc)
            res.push_back(*item);
        return res;
    }

    std::vector< Eigen::ArrayXXd  >  set(const  std::shared_ptr<StOpt::PayOffBase>   &p_funcValue, const Eigen::ArrayXXd &p_particles) const
    {
        std::vector< std::shared_ptr< Eigen::ArrayXXd > > res = StOpt::FinalStepDP::operator()(p_funcValue->getFunction(), p_particles);
        std::vector<Eigen::ArrayXXd >result;
        result.reserve(res.size());
        for (const auto &item : res)
            result.push_back(*item);
        return result;
    }
};

// wrapper for Transition step for tree
//*************************************

class PyTransitionStepTreeDP: StOpt::TransitionStepTreeDP
{

public :
    PyTransitionStepTreeDP(const std::shared_ptr< StOpt::SpaceGrid >   &p_gridCurrent,  const std::shared_ptr< StOpt::SpaceGrid >   &p_gridPrevious, const  std::shared_ptr< StOpt::OptimizerDPTreeBase>   &p_optimizer):
        StOpt::TransitionStepTreeDP(std::dynamic_pointer_cast<StOpt::FullGrid>(p_gridCurrent),
                                    std::dynamic_pointer_cast<StOpt::FullGrid>(p_gridPrevious),
                                    p_optimizer
#ifdef USE_MPI
                                    , boost::mpi::communicator()
#endif
                                   ) {  }

    py::list  oneStepWrap(const std::vector< Eigen::ArrayXXd >  &p_phiIn,
                          const std::shared_ptr< StOpt::Tree>     &p_condExp) const
    {
        std::vector< std::shared_ptr< Eigen::ArrayXXd > >  input;
        input.reserve(p_phiIn.size());
        for (const auto &item : p_phiIn)
        {
            input.push_back(std::make_shared<Eigen::ArrayXXd>(item));
        }
        std::pair< std::vector< std::shared_ptr< Eigen::ArrayXXd > >, std::vector<  std::shared_ptr<  Eigen::ArrayXXd > > > retLoc = StOpt::TransitionStepTreeDP::oneStep(input, p_condExp);
        py::list res ;
        py::list res1;
        for (const auto &item : retLoc.first)
            res1.append(*item);
        py::list res2;
        for (const auto &item : retLoc.second)
            res2.append(*item);
        res.append(res1);
        res.append(res2);
        return res;
    }

    void dumpContinuationValuesWrap(std::shared_ptr<BinaryFileArchiveStOpt>  p_ar, const std::string &p_name, const int &p_iStep, const  py::list  &p_phiIn,  const  py::list   &p_control, const   std::shared_ptr<StOpt::Tree>    &p_condExp) const
    {
        std::vector< std::shared_ptr< Eigen::ArrayXXd > >  phiIn = convertFromListShPtr<Eigen::ArrayXXd>(p_phiIn);
        std::vector< std::shared_ptr< Eigen::ArrayXXd > >  control = convertFromListShPtr<Eigen::ArrayXXd>(p_control);
        StOpt::TransitionStepTreeDP::dumpContinuationValues(p_ar, p_name, p_iStep, phiIn, control, p_condExp);
    }
};


// Wrapper for SimulateStepTree
class PySimulateStepTree : StOpt::SimulateStepTree

{
public :
    PySimulateStepTree(BinaryFileArchiveStOpt &p_ar,  const int &p_iStep,  const std::string &p_nameCont,
                       const   std::shared_ptr< StOpt::SpaceGrid > &p_pGridFollowing,
                       const  std::shared_ptr< StOpt::OptimizerDPTreeBase > &p_pOptimize):
        StOpt::SimulateStepTree(p_ar, p_iStep, p_nameCont, p_pGridFollowing, p_pOptimize
#ifdef USE_MPI
                                , boost::mpi::communicator()
#endif
                               )
    {}

    py::list  oneStepWrap(const py::list    &p_statevector,  const Eigen::ArrayXXd   &p_phiInOut) const
    {
        std::vector<StOpt::StateTreeStocks> stateLoc = convertFromList<StOpt::StateTreeStocks>(p_statevector);
        Eigen::ArrayXXd phiLoc(p_phiInOut);
        StOpt::SimulateStepTree::oneStep(stateLoc, phiLoc);
        py::list ret;
        ret.append(py::cast(stateLoc));
        ret.append(py::cast(phiLoc));
        return ret ;
    }
};


// Wrapper for SimulateStepTreeControl
class PySimulateStepTreeControl : StOpt::SimulateStepTreeControl

{
public :
    PySimulateStepTreeControl(BinaryFileArchiveStOpt &p_ar,  const int &p_iStep,  const std::string &p_nameCont,
                              const   std::shared_ptr< StOpt::SpaceGrid > &p_pGridFollowing, const  std::shared_ptr< StOpt::OptimizerDPTreeBase > &p_pOptimize):
        StOpt::SimulateStepTreeControl(p_ar, p_iStep, p_nameCont, p_pGridFollowing, p_pOptimize
#ifdef USE_MPI
                                       , boost::mpi::communicator()
#endif
                                      )
    {}

    py::list  oneStepWrap(const py::list    &p_statevector,  const Eigen::ArrayXXd   &p_phiInOut) const
    {
        std::vector<StOpt::StateTreeStocks> stateLoc = convertFromList<StOpt::StateTreeStocks>(p_statevector);
        Eigen::ArrayXXd phiLoc(p_phiInOut);
        StOpt::SimulateStepTreeControl::oneStep(stateLoc, phiLoc);
        py::list ret;
        ret.append(py::cast(stateLoc));
        ret.append(py::cast(phiLoc));
        return ret ;
    }
};


/// mapping for transition step with regressions
class PyTransitionStepRegressionDP: StOpt::TransitionStepRegressionDP
{
public :
    PyTransitionStepRegressionDP(const std::shared_ptr< StOpt::SpaceGrid >   &p_gridCurrent,  const std::shared_ptr< StOpt::SpaceGrid >   &p_gridPrevious, const  std::shared_ptr< StOpt::OptimizerDPBase>   &p_optimizer):
        StOpt::TransitionStepRegressionDP(std::dynamic_pointer_cast<StOpt::FullGrid>(p_gridCurrent),
                                          std::dynamic_pointer_cast<StOpt::FullGrid>(p_gridPrevious),
                                          p_optimizer
#ifdef USE_MPI
                                          , boost::mpi::communicator()
#endif
                                         ) {  }

    py::list  oneStepWrap(const std::vector< Eigen::ArrayXXd >  &p_phiIn,
                          const std::shared_ptr< StOpt::BaseRegression>     &p_condExp) const
    {
        std::vector< std::shared_ptr< Eigen::ArrayXXd > >  input;
        input.reserve(p_phiIn.size());
        for (const auto &item : p_phiIn)
        {
            input.push_back(std::make_shared<Eigen::ArrayXXd>(item));
        }
        std::pair< std::vector< std::shared_ptr< Eigen::ArrayXXd > >, std::vector<  std::shared_ptr<  Eigen::ArrayXXd > > > retLoc = StOpt::TransitionStepRegressionDP::oneStep(input, p_condExp);
        py::list res ;
        py::list res1;
        for (const auto &item : retLoc.first)
            res1.append(*item);
        py::list res2;
        for (const auto &item : retLoc.second)
            res2.append(*item);
        res.append(res1);
        res.append(res2);
        return res;
    }

    void dumpContinuationValuesWrap(std::shared_ptr<BinaryFileArchiveStOpt>  p_ar, const std::string &p_name, const int &p_iStep, const  py::list  &p_phiIn,  const  py::list   &p_control, const   std::shared_ptr<StOpt::BaseRegression>    &p_condExp) const
    {
        std::vector< std::shared_ptr< Eigen::ArrayXXd > >  phiIn = convertFromListShPtr<Eigen::ArrayXXd>(p_phiIn);
        std::vector< std::shared_ptr< Eigen::ArrayXXd > >  control = convertFromListShPtr<Eigen::ArrayXXd>(p_control);
        StOpt::TransitionStepRegressionDP::dumpContinuationValues(p_ar, p_name, p_iStep, phiIn, control, p_condExp);
    }

    void dumpBellmanValuesWrap(std::shared_ptr<BinaryFileArchiveStOpt>  p_ar, const std::string &p_name, const int &p_iStep, const  py::list  &p_phiIn,   const   std::shared_ptr<StOpt::BaseRegression>    &p_condExp) const
    {
        std::vector< std::shared_ptr< Eigen::ArrayXXd > >  phiIn = convertFromListShPtr<Eigen::ArrayXXd>(p_phiIn);
        StOpt::TransitionStepRegressionDP::dumpBellmanValues(p_ar, p_name, p_iStep, phiIn,  p_condExp);
    }
};


// Wrapper for SimulateStepRegression
class PySimulateStepRegression : StOpt::SimulateStepRegression

{
public :
    PySimulateStepRegression(BinaryFileArchiveStOpt &p_ar,  const int &p_iStep,  const std::string &p_nameCont,
                             const   std::shared_ptr< StOpt::SpaceGrid > &p_pGridFollowing,
                             const  std::shared_ptr< StOpt::OptimizerDPBase > &p_pOptimize):
        StOpt::SimulateStepRegression(p_ar, p_iStep, p_nameCont, p_pGridFollowing, p_pOptimize
#ifdef USE_MPI
                                      , boost::mpi::communicator()
#endif
                                     )
    {}

    py::list  oneStepWrap(const py::list    &p_statevector,  const Eigen::ArrayXXd   &p_phiInOut) const
    {
        std::vector<StOpt::StateWithStocks> stateLoc = convertFromList<StOpt::StateWithStocks>(p_statevector);
        Eigen::ArrayXXd phiLoc(p_phiInOut);
        StOpt::SimulateStepRegression::oneStep(stateLoc, phiLoc);
        py::list ret;
        ret.append(py::cast(stateLoc));
        ret.append(py::cast(phiLoc));
        return ret ;
    }
};


// Wrapper for SimulateStepRegressionControl
class PySimulateStepRegressionControl : StOpt::SimulateStepRegressionControl
{
public :
    PySimulateStepRegressionControl(BinaryFileArchiveStOpt &p_ar,  const int &p_iStep,  const std::string &p_nameCont,
                                    const   std::shared_ptr< StOpt::SpaceGrid > &p_pGridFollowing, const  std::shared_ptr< StOpt::OptimizerBaseInterp > &p_pOptimize):
        StOpt::SimulateStepRegressionControl(p_ar, p_iStep, p_nameCont, p_pGridFollowing, p_pOptimize
#ifdef USE_MPI
                                             , boost::mpi::communicator()
#endif
                                            )
    {}

    py::list  oneStepWrap(const py::list    &p_statevector,  const Eigen::ArrayXXd   &p_phiInOut) const
    {
        std::vector<StOpt::StateWithStocks> stateLoc = convertFromList<StOpt::StateWithStocks>(p_statevector);
        Eigen::ArrayXXd phiLoc(p_phiInOut);
        StOpt::SimulateStepRegressionControl::oneStep(stateLoc, phiLoc);
        py::list ret;
        ret.append(py::cast(stateLoc));
        ret.append(py::cast(phiLoc));
        return ret ;
    }
};


/// mapping for transition step with switching
class PyTransitionStepRegressionSwitch: StOpt::TransitionStepRegressionSwitch
{
public :
    PyTransitionStepRegressionSwitch(const  std::vector< std::shared_ptr<  StOpt::RegularSpaceIntGrid> >  &p_gridCurrent,
                                     const  std::vector< std::shared_ptr<StOpt::RegularSpaceIntGrid> >  &p_gridPrevious,
                                     const  std::shared_ptr<StOpt::OptimizerSwitchBase > &p_optimizer):
        StOpt::TransitionStepRegressionSwitch(p_gridCurrent, p_gridPrevious, p_optimizer
#ifdef USE_MPI
                                              , boost::mpi::communicator()
#endif
                                             ) {  }


    py::list  oneStepWrap(const std::vector< Eigen::ArrayXXd >  &p_phiIn,
                          const std::shared_ptr< StOpt::BaseRegression>     &p_condExp) const
    {
        std::vector< std::shared_ptr< Eigen::ArrayXXd > >  input;
        input.reserve(p_phiIn.size());
        for (const auto &item : p_phiIn)
        {
            input.push_back(std::make_shared<Eigen::ArrayXXd>(item));
        }
        std::vector< std::shared_ptr< Eigen::ArrayXXd > > retLoc = StOpt::TransitionStepRegressionSwitch::oneStep(input, p_condExp);
        py::list res;
        for (const auto &item : retLoc)
            res.append(*item);
        return res;
    }

    void dumpContinuationValuesWrap(std::shared_ptr<BinaryFileArchiveStOpt>  p_ar, const std::string &p_name, const int &p_iStep, const  py::list  &p_phiIn,  const   std::shared_ptr<StOpt::BaseRegression>    &p_condExp) const
    {
        std::vector< std::shared_ptr< Eigen::ArrayXXd > >  phiIn = convertFromListShPtr<Eigen::ArrayXXd>(p_phiIn);
        StOpt::TransitionStepRegressionSwitch::dumpContinuationValues(p_ar, p_name, p_iStep, phiIn, p_condExp);
    }

};


// Wrapper for SimulateStepSwitch
class PySimulateStepSwitch : StOpt::SimulateStepSwitch

{
public :
    PySimulateStepSwitch(BinaryFileArchiveStOpt &p_ar,  const int &p_iStep,  const std::string &p_nameCont,
                         const   std::vector<std::shared_ptr< StOpt::RegularSpaceIntGrid > > &p_pGridFollowing,
                         const   std::shared_ptr< StOpt::OptimizerSwitchBase > &p_pOptimize):
        StOpt::SimulateStepSwitch(p_ar, p_iStep, p_nameCont, p_pGridFollowing, p_pOptimize
#ifdef USE_MPI
                                  , boost::mpi::communicator()
#endif
                                 )
    {}

    py::list  oneStepWrap(const py::list    &p_statevector,  const Eigen::ArrayXXd   &p_phiInOut) const
    {
        std::vector<StOpt::StateWithIntState> stateLoc = convertFromList<StOpt::StateWithIntState>(p_statevector);
        Eigen::ArrayXXd phiLoc(p_phiInOut);
        StOpt::SimulateStepSwitch::oneStep(stateLoc, phiLoc);
        py::list ret;
        ret.append(py::cast(stateLoc));
        ret.append(py::cast(phiLoc));
        return ret ;
    }
};


#ifdef USE_MPI

class PyFinalStepDPDist: StOpt::FinalStepDPDist
{
public:

    PyFinalStepDPDist(const  std::shared_ptr<StOpt::SpaceGrid> &p_pGridCurrent, const int &p_nbRegime,
                      const Eigen::Array< bool, Eigen::Dynamic, 1>   &p_bdimToSplit):
        StOpt::FinalStepDPDist(std::dynamic_pointer_cast<StOpt::FullGrid>(p_pGridCurrent), p_nbRegime, p_bdimToSplit, boost::mpi::communicator()) {}

    std::vector< Eigen::ArrayXXd >   operator()(py::object &p_funcValue, const Eigen::ArrayXXd &p_particles) const
    {
        auto lambda = [p_funcValue](const int &p_i, const Eigen::ArrayXd & p_x, const Eigen::ArrayXd & p_y)->double { return p_funcValue(p_i, p_x, p_y).cast<double>();};

        std::vector< std::shared_ptr<Eigen::ArrayXXd>  >  resLoc = StOpt::FinalStepDPDist::operator()(std::function<double(const int &, const Eigen::ArrayXd &, const Eigen::ArrayXd &)> (lambda), p_particles);
        std::vector< Eigen::ArrayXXd > res;
        res.reserve(resLoc.size());
        for (const auto   &item : resLoc)
            res.push_back(*item);
        return res;
    }
    std::vector< Eigen::ArrayXXd  >  set(const  std::shared_ptr<StOpt::PayOffBase>   &p_funcValue, const Eigen::ArrayXXd &p_particles) const
    {
        std::vector< std::shared_ptr< Eigen::ArrayXXd > > res = StOpt::FinalStepDPDist::operator()(p_funcValue->getFunction(), p_particles);
        std::vector<Eigen::ArrayXXd >result;
        result.reserve(res.size());
        for (const auto &item : res)
            result.push_back(*item);
        return result;
    }
};

class PyFinalStepZeroDist: StOpt::FinalStepZeroDist<StOpt::RegularSpaceIntGrid>
{
public:

    PyFinalStepZeroDist(const  std::vector< std::shared_ptr< StOpt::RegularSpaceIntGrid> >   &p_pGridCurrent,
                        const  std::vector< Eigen::Array< bool, Eigen::Dynamic, 1> > &p_bdimToSplit):
        StOpt::FinalStepZeroDist<StOpt::RegularSpaceIntGrid>(p_pGridCurrent, p_bdimToSplit, boost::mpi::communicator()) {}

    std::vector< Eigen::ArrayXXd >   operator()(const int   &p_nbSimul) const
    {
        std::vector< std::shared_ptr<Eigen::ArrayXXd>  >  resLoc = StOpt::FinalStepZeroDist<StOpt::RegularSpaceIntGrid>::operator()(p_nbSimul);
        std::vector< Eigen::ArrayXXd > res;
        res.reserve(resLoc.size());
        for (const auto   &item : resLoc)
            res.push_back(*item);
        return res;
    }
};

class PyTransitionStepRegressionDPDist: StOpt::TransitionStepRegressionDPDist
{
public :
    PyTransitionStepRegressionDPDist(const std::shared_ptr< StOpt::SpaceGrid >   &p_gridCurrent,  const std::shared_ptr< StOpt::SpaceGrid >   &p_gridPrevious, const  std::shared_ptr< StOpt::OptimizerDPBase>   &p_optimizer):
        StOpt::TransitionStepRegressionDPDist(std::dynamic_pointer_cast<StOpt::FullGrid>(p_gridCurrent),
                                              std::dynamic_pointer_cast<StOpt::FullGrid>(p_gridPrevious),
                                              p_optimizer, boost::mpi::communicator()) {  }

    py::list  oneStepWrap(const std::vector< Eigen::ArrayXXd >  &p_phiIn,
                          const std::shared_ptr< StOpt::BaseRegression>     &p_condExp) const
    {
        std::vector< std::shared_ptr< Eigen::ArrayXXd > >  input;
        input.reserve(p_phiIn.size());
        for (const auto &item : p_phiIn)
        {
            input.push_back(std::make_shared<Eigen::ArrayXXd>(item));
        }
        std::pair< std::vector< std::shared_ptr< Eigen::ArrayXXd > >, std::vector<  std::shared_ptr<  Eigen::ArrayXXd > > > retLoc = StOpt::TransitionStepRegressionDPDist::oneStep(input, p_condExp);
        py::list res ;
        py::list res1;
        for (const auto &item : retLoc.first)
            res1.append(*item);
        py::list res2;
        for (const auto &item : retLoc.second)
            res2.append(*item);
        res.append(res1);
        res.append(res2);
        return res;
    }

    void dumpContinuationValuesWrap(std::shared_ptr<BinaryFileArchiveStOpt>  p_ar, const std::string &p_name, const int &p_iStep, const  py::list  &p_phiIn,  const  py::list   &p_control, const   std::shared_ptr<StOpt::BaseRegression>    &p_condExp, const bool &p_bOneFile) const
    {
        std::vector< std::shared_ptr< Eigen::ArrayXXd > >  phiIn = convertFromListShPtr<Eigen::ArrayXXd>(p_phiIn);
        std::vector< std::shared_ptr< Eigen::ArrayXXd > >  control = convertFromListShPtr<Eigen::ArrayXXd>(p_control);
        StOpt::TransitionStepRegressionDPDist::dumpContinuationValues(p_ar, p_name, p_iStep, phiIn, control, p_condExp, p_bOneFile);
    }

    void dumpBellmanValuesWrap(std::shared_ptr<BinaryFileArchiveStOpt>  p_ar, const std::string &p_name, const int &p_iStep, const  py::list  &p_phiIn,   const   std::shared_ptr<StOpt::BaseRegression>    &p_condExp, const bool &p_bOneFile) const
    {
        std::vector< std::shared_ptr< Eigen::ArrayXXd > >  phiIn = convertFromListShPtr<Eigen::ArrayXXd>(p_phiIn);
        StOpt::TransitionStepRegressionDPDist::dumpBellmanValues(p_ar, p_name, p_iStep, phiIn,  p_condExp, p_bOneFile);
    }
};


// Wrapper for SimulateStepRegression
class PySimulateStepRegressionDist : StOpt::SimulateStepRegressionDist

{
public :
    PySimulateStepRegressionDist(BinaryFileArchiveStOpt &p_ar,  const int &p_iStep,  const std::string &p_nameCont,
                                 const   std::shared_ptr< StOpt::SpaceGrid > &p_pGridFollowing,
                                 const  std::shared_ptr< StOpt::OptimizerDPBase > &p_pOptimize,
                                 const bool &p_bOneFile):
        StOpt::SimulateStepRegressionDist(p_ar, p_iStep, p_nameCont, std::dynamic_pointer_cast<StOpt::FullGrid>(p_pGridFollowing), p_pOptimize, p_bOneFile, boost::mpi::communicator())
    {}

    py::list  oneStepWrap(const py::list    &p_statevector,  const Eigen::ArrayXXd   &p_phiInOut) const
    {
        std::vector<StOpt::StateWithStocks> stateLoc = convertFromList<StOpt::StateWithStocks>(p_statevector);
        Eigen::ArrayXXd phiLoc(p_phiInOut);
        StOpt::SimulateStepRegressionDist::oneStep(stateLoc, phiLoc);
        py::list ret;
        ret.append(py::cast(stateLoc));
        ret.append(py::cast(phiLoc));
        return ret ;
    }
};


// Wrapper for SimulateStepRegressionControl
class PySimulateStepRegressionControlDist : StOpt::SimulateStepRegressionControlDist

{
public :
    PySimulateStepRegressionControlDist(BinaryFileArchiveStOpt &p_ar,  const int &p_iStep,  const std::string &p_nameCont,
                                        const   std::shared_ptr< StOpt::SpaceGrid > &p_pGridCurrent,
                                        const   std::shared_ptr< StOpt::SpaceGrid > &p_pGridFollowing,
                                        const  std::shared_ptr< StOpt::OptimizerDPBase > &p_pOptimize,
                                        const bool &p_bOneFile):
        StOpt::SimulateStepRegressionControlDist(p_ar, p_iStep, p_nameCont, std::dynamic_pointer_cast<StOpt::FullGrid>(p_pGridCurrent),
                std::dynamic_pointer_cast<StOpt::FullGrid>(p_pGridFollowing), p_pOptimize, p_bOneFile, boost::mpi::communicator())
    {}

    py::list  oneStepWrap(const py::list    &p_statevector,  const Eigen::ArrayXXd   &p_phiInOut) const
    {
        std::vector<StOpt::StateWithStocks> stateLoc = convertFromList<StOpt::StateWithStocks>(p_statevector);
        Eigen::ArrayXXd phiLoc(p_phiInOut);
        StOpt::SimulateStepRegressionControlDist::oneStep(stateLoc, phiLoc);
        py::list ret;
        ret.append(py::cast(stateLoc));
        ret.append(py::cast(phiLoc));
        return ret ;
    }
};

class PyTransitionStepTreeDPDist: StOpt::TransitionStepTreeDPDist
{
public :
    PyTransitionStepTreeDPDist(const std::shared_ptr< StOpt::SpaceGrid >   &p_gridCurrent,  const std::shared_ptr< StOpt::SpaceGrid >   &p_gridPrevious, const  std::shared_ptr< StOpt::OptimizerDPTreeBase>   &p_optimizer):
        StOpt::TransitionStepTreeDPDist(std::dynamic_pointer_cast<StOpt::FullGrid>(p_gridCurrent),
                                        std::dynamic_pointer_cast<StOpt::FullGrid>(p_gridPrevious),
                                        p_optimizer, boost::mpi::communicator()) {  }

    py::list  oneStepWrap(const std::vector< Eigen::ArrayXXd >  &p_phiIn,
                          const std::shared_ptr< StOpt::Tree>     &p_condExp) const
    {
        std::vector< std::shared_ptr< Eigen::ArrayXXd > >  input;
        input.reserve(p_phiIn.size());
        for (const auto &item : p_phiIn)
        {
            input.push_back(std::make_shared<Eigen::ArrayXXd>(item));
        }
        std::pair< std::vector< std::shared_ptr< Eigen::ArrayXXd > >, std::vector<  std::shared_ptr<  Eigen::ArrayXXd > > > retLoc = StOpt::TransitionStepTreeDPDist::oneStep(input, p_condExp);
        py::list res ;
        py::list res1;
        for (const auto &item : retLoc.first)
            res1.append(*item);
        py::list res2;
        for (const auto &item : retLoc.second)
            res2.append(*item);
        res.append(res1);
        res.append(res2);
        return res;
    }

    void dumpContinuationValuesWrap(std::shared_ptr<BinaryFileArchiveStOpt>  p_ar, const std::string &p_name, const int &p_iStep, const  py::list  &p_phiIn,  const  py::list   &p_control, const   std::shared_ptr<StOpt::Tree>    &p_condExp, const bool &p_bOneFile) const
    {
        std::vector< std::shared_ptr< Eigen::ArrayXXd > >  phiIn = convertFromListShPtr<Eigen::ArrayXXd>(p_phiIn);
        std::vector< std::shared_ptr< Eigen::ArrayXXd > >  control = convertFromListShPtr<Eigen::ArrayXXd>(p_control);
        StOpt::TransitionStepTreeDPDist::dumpContinuationValues(p_ar, p_name, p_iStep, phiIn, control, p_condExp, p_bOneFile);
    }
};


// Wrapper for SimulateStepTree
class PySimulateStepTreeDist : StOpt::SimulateStepTreeDist

{
public :
    PySimulateStepTreeDist(BinaryFileArchiveStOpt &p_ar,  const int &p_iStep,  const std::string &p_nameCont,
                           const   std::shared_ptr< StOpt::SpaceGrid > &p_pGridFollowing,
                           const  std::shared_ptr< StOpt::OptimizerDPTreeBase > &p_pOptimize,
                           const bool &p_bOneFile):
        StOpt::SimulateStepTreeDist(p_ar, p_iStep, p_nameCont, std::dynamic_pointer_cast<StOpt::FullGrid>(p_pGridFollowing), p_pOptimize, p_bOneFile, boost::mpi::communicator())
    {}

    py::list  oneStepWrap(const py::list    &p_statevector,  const Eigen::ArrayXXd   &p_phiInOut) const
    {
        std::vector<StOpt::StateTreeStocks> stateLoc = convertFromList<StOpt::StateTreeStocks>(p_statevector);
        Eigen::ArrayXXd phiLoc(p_phiInOut);
        StOpt::SimulateStepTreeDist::oneStep(stateLoc, phiLoc);
        py::list ret;
        ret.append(py::cast(stateLoc));
        ret.append(py::cast(phiLoc));
        return ret ;
    }
};


// Wrapper for SimulateStepTreeControl
class PySimulateStepTreeControlDist : StOpt::SimulateStepTreeControlDist

{
public :
    PySimulateStepTreeControlDist(BinaryFileArchiveStOpt &p_ar,  const int &p_iStep,  const std::string &p_nameCont,
                                  const   std::shared_ptr< StOpt::SpaceGrid > &p_pGridCurrent,
                                  const   std::shared_ptr< StOpt::SpaceGrid > &p_pGridFollowing,
                                  const  std::shared_ptr< StOpt::OptimizerDPTreeBase > &p_pOptimize,
                                  const bool &p_bOneFile):
        StOpt::SimulateStepTreeControlDist(p_ar, p_iStep, p_nameCont, std::dynamic_pointer_cast<StOpt::FullGrid>(p_pGridCurrent),
                                           std::dynamic_pointer_cast<StOpt::FullGrid>(p_pGridFollowing), p_pOptimize, p_bOneFile, boost::mpi::communicator())
    {}

    py::list  oneStepWrap(const py::list    &p_statevector,  const Eigen::ArrayXXd   &p_phiInOut) const
    {
        std::vector<StOpt::StateTreeStocks> stateLoc = convertFromList<StOpt::StateTreeStocks>(p_statevector);
        Eigen::ArrayXXd phiLoc(p_phiInOut);
        StOpt::SimulateStepTreeControlDist::oneStep(stateLoc, phiLoc);
        py::list ret;
        ret.append(py::cast(stateLoc));
        ret.append(py::cast(phiLoc));
        return ret ;
    }
};


/// mapping for transition step with switching
class PyTransitionStepRegressionSwitchDist: StOpt::TransitionStepRegressionSwitchDist
{
public :
    PyTransitionStepRegressionSwitchDist(const  std::vector< std::shared_ptr<  StOpt::RegularSpaceIntGrid> >  &p_gridCurrent,
                                         const  std::vector< std::shared_ptr<StOpt::RegularSpaceIntGrid> >  &p_gridPrevious,
                                         const  std::shared_ptr<StOpt::OptimizerSwitchBase > &p_optimizer):
        StOpt::TransitionStepRegressionSwitchDist(p_gridCurrent, p_gridPrevious, p_optimizer, boost::mpi::communicator()) {  }


    py::list  oneStepWrap(const std::vector< Eigen::ArrayXXd >  &p_phiIn,
                          const std::shared_ptr< StOpt::BaseRegression>     &p_condExp) const
    {
        std::vector< std::shared_ptr< Eigen::ArrayXXd > >  input;
        input.reserve(p_phiIn.size());
        for (const auto &item : p_phiIn)
        {
            input.push_back(std::make_shared<Eigen::ArrayXXd>(item));
        }
        std::vector< std::shared_ptr< Eigen::ArrayXXd > > retLoc = StOpt::TransitionStepRegressionSwitchDist::oneStep(input, p_condExp);
        py::list res;
        for (const auto &item : retLoc)
            res.append(*item);
        return res;
    }

    void dumpContinuationValuesWrap(std::shared_ptr<BinaryFileArchiveStOpt>  p_ar, const std::string &p_name, const int &p_iStep, const  py::list  &p_phiIn,  const   std::shared_ptr<StOpt::BaseRegression>    &p_condExp) const
    {
        std::vector< std::shared_ptr< Eigen::ArrayXXd > >  phiIn = convertFromListShPtr<Eigen::ArrayXXd>(p_phiIn);
        StOpt::TransitionStepRegressionSwitchDist::dumpContinuationValues(p_ar, p_name, p_iStep, phiIn, p_condExp);
    }

};



Eigen::ArrayXd  pyReconstructProc0Mpi(const Eigen::ArrayXd &p_point, const std::shared_ptr< StOpt::SpaceGrid> &p_grid, const Eigen::ArrayXXd    &p_values, const Eigen::Array< bool, Eigen::Dynamic, 1>   &p_bdimToSplit)
{
    std::shared_ptr< Eigen::ArrayXXd > values = std::make_shared<Eigen::ArrayXXd > (p_values);
    return reconstructProc0Mpi(p_point, std::dynamic_pointer_cast<StOpt::FullGrid>(p_grid), values, p_bdimToSplit, boost::mpi::communicator());
}


Eigen::ArrayXd  pyReconstructProc0ForIntMpi(const Eigen::ArrayXi &p_point, const std::shared_ptr< StOpt::RegularSpaceIntGrid> &p_grid, const Eigen::ArrayXXd    &p_values, const Eigen::Array< bool, Eigen::Dynamic, 1>   &p_bdimToSplit)
{
    std::shared_ptr< Eigen::ArrayXXd > values = std::make_shared<Eigen::ArrayXXd > (p_values);
    return reconstructProc0ForIntMpi(p_point, p_grid, values, p_bdimToSplit, boost::mpi::communicator());
}

#endif



namespace py = pybind11;

/// \brief Encapsulation for regression module
PYBIND11_MODULE(StOptGlobal, m)
{


    py::class_<StOpt::StateWithStocks>(m, "StateWithStocks")
    .def(py::init< const int &, const Eigen::ArrayXd &, const Eigen::ArrayXd &>())
    .def("getPtStock", &StOpt::StateWithStocks::getPtStock, py::return_value_policy::copy) // copy instead fo reference
    .def("getRegime", &StOpt::StateWithStocks::getRegime)
    .def("setPtStock", &StOpt::StateWithStocks::setPtStock)
    .def("getStochasticRealization", &StOpt::StateWithStocks::getStochasticRealization,  py::return_value_policy::copy)
    .def("setPtOneStock", &StOpt::StateWithStocks::setPtOneStock)
    .def("setStochasticRealization", &StOpt::StateWithStocks::setStochasticRealization)
    .def("setRegime", &StOpt::StateWithStocks::setRegime)
    ;

    py::class_<StOpt::StateWithIntState>(m, "StateWithIntState")
    .def(py::init< const int &, const Eigen::ArrayXi &, const Eigen::ArrayXd &>())
    .def("getPtState", &StOpt::StateWithIntState::getPtState, py::return_value_policy::copy) // copy instead fo reference
    .def("getRegime", &StOpt::StateWithIntState::getRegime)
    .def("setPtState", &StOpt::StateWithIntState::setPtState)
    .def("getStochasticRealization", &StOpt::StateWithIntState::getStochasticRealization,  py::return_value_policy::copy)
    .def("setPtOneState", &StOpt::StateWithIntState::setPtOneState)
    .def("setStochasticRealization", &StOpt::StateWithIntState::setStochasticRealization)
    .def("setRegime", &StOpt::StateWithIntState::setRegime)
    ;

    py::class_<StOpt::StateTreeStocks>(m, "StateTreeStocks")
    .def(py::init< const int &, const Eigen::ArrayXd &, const int &>())
    .def("getPtStock", &StOpt::StateTreeStocks::getPtStock, py::return_value_policy::copy) // copy instead fo reference
    .def("getRegime", &StOpt::StateTreeStocks::getRegime)
    .def("setPtStock", &StOpt::StateTreeStocks::setPtStock)
    .def("getStochasticRealization", &StOpt::StateTreeStocks::getStochasticRealization,  py::return_value_policy::copy)
    .def("setStochasticRealization", &StOpt::StateTreeStocks::setStochasticRealization)
    .def("setRegime", &StOpt::StateTreeStocks::setRegime)
    .def("setPtOneStock", &StOpt::StateTreeStocks::setPtOneStock)
    .def("getPtOneStock", &StOpt::StateTreeStocks::getPtOneStock)
    ;


    // map simulator DP base
    //**********************
    py::class_<StOpt::SimulatorDPBase, std::shared_ptr<StOpt::SimulatorDPBase>, PySimulatorDPBase >(m, "SimulatorDPBase")
    .def(py::init<>())
    .def("getParticles", &StOpt::SimulatorDPBase::getParticles)
    .def("stepForward", &StOpt::SimulatorDPBase::stepForward)
    .def("stepBackward", &StOpt::SimulatorDPBase::stepBackward)
    .def("stepForwardAndGetParticles", &StOpt::SimulatorDPBase::stepForwardAndGetParticles)
    .def("stepBackwardAndGetParticles", &StOpt::SimulatorDPBase::stepBackwardAndGetParticles)
    .def("getDimension", &StOpt::SimulatorDPBase::getDimension)
    .def("getNbStep", &StOpt::SimulatorDPBase::getNbStep)
    .def("getStep", &StOpt::SimulatorDPBase::getStep)
    .def("getCurrentStep", &StOpt::SimulatorDPBase::getCurrentStep)
    .def("getNbSimul", &StOpt::SimulatorDPBase::getNbSimul)
    .def("getActuStep", &StOpt::SimulatorDPBase::getActuStep)
    .def("getActu", &StOpt::SimulatorDPBase::getActu)
    ;

    // map simulator tree DP base
    //***************************
    py::class_<StOpt::SimulatorDPBaseTree, std::shared_ptr<StOpt::SimulatorDPBaseTree>, PySimulatorDPBaseTree >(m, "SimulatorDPBaseTree")
    .def(py::init<>())
    .def("getNodeReachedInForward", &StOpt::SimulatorDPBaseTree::getNodeReachedInForward)
    .def("stepBackward", &StOpt::SimulatorDPBaseTree::stepBackward)
    .def("stepForward", &StOpt::SimulatorDPBaseTree::stepForward)
    .def("getDimension", &StOpt::SimulatorDPBaseTree::getDimension)
    .def("getNbStep", &StOpt::SimulatorDPBaseTree::getNbStep)
    .def("getNbNodes", &StOpt::SimulatorDPBaseTree::getNbNodes)
    .def("getNbNodesNext", &StOpt::SimulatorDPBaseTree::getNbNodesNext)
    .def("getDates", &StOpt::SimulatorDPBaseTree::getDates)
    .def("getBackLastDateIndex", &StOpt::SimulatorDPBaseTree::getBackLastDateIndex)
    .def("getConnected", &StOpt::SimulatorDPBaseTree::getConnected)
    .def("getProba", &StOpt::SimulatorDPBaseTree::getProba)
    .def("getNodes", &StOpt::SimulatorDPBaseTree::getNodes)
    .def("getNodesNext", &StOpt::SimulatorDPBaseTree::getNodesNext)
    .def("getNbSimul", &StOpt::SimulatorDPBaseTree::getNbSimul)
    .def("getValueAssociatedToNode", &StOpt::SimulatorDPBaseTree::getValueAssociatedToNode)
    .def("getNodeAssociatedToSim", &StOpt::SimulatorDPBaseTree::getNodeAssociatedToSim)
    ;

    // map optimizer DP base
    //**********************
    py::class_<StOpt::OptimizerDPBase, std::shared_ptr<StOpt::OptimizerDPBase>, PyOptimizerDPBase>(m, "OptimizerDPBase")
    .def("getCone",  [](const StOpt::OptimizerDPBase & p_x, const  std::vector< Eigen::ArrayXd >    &p_region)
    {
        std::vector< std::array< double, 2> > region;
        region.reserve(p_region.size());
        for (const auto &item : p_region)
        {
            region.push_back(std::array< double, 2> {{ item[0], item[1]}});
        }
        std::vector< std::array< double, 2> > resLoc = p_x.getCone(region);
        std::vector< Eigen::ArrayXd > result;
        result.reserve(resLoc.size());
        for (const auto &item : resLoc)
        {
            Eigen::ArrayXd dom(2);
            dom(0) = item[0];
            dom(1) = item[1];
            result.push_back(dom);
        }
        return result;
    }
        )
    .def("getDimensionToSplit",  &StOpt::OptimizerDPBase::getDimensionToSplit)
    .def("stepOptimize", [](const StOpt::OptimizerDPBase & p_x,
                            const std::shared_ptr< StOpt::SpaceGrid> &p_grid,
                            const Eigen::ArrayXd   & p_reg,
                            const std::vector<StOpt::ContinuationValue> &p_cont,
                            const std::vector < Eigen::ArrayXXd  > &p_val)
    {
        std::vector < std::shared_ptr< Eigen::ArrayXXd > > val;
        val.reserve(p_val.size());
        for (const auto &item : p_val)
            val.push_back(std::make_shared<Eigen::ArrayXXd>(item));
        std::pair< Eigen::ArrayXXd, Eigen::ArrayXXd>  resLoc = p_x.stepOptimize(p_grid, p_reg, p_cont, val);
        py::list res ;
        res.append(py::cast(resLoc.first));
        res.append(py::cast(resLoc.second));
        return res;
    }
        )
    .def("stepSimulate", &StOpt::OptimizerDPBase::stepSimulate)
    .def("stepSimulateControl", &StOpt::OptimizerDPBase::stepSimulateControl)
    .def("getNbRegime", &StOpt::OptimizerDPBase::getNbRegime)
    .def("getNbControl", &StOpt::OptimizerDPBase::getNbControl)
    .def("getSimulator", &StOpt::OptimizerDPBase::getSimulator)
    .def("getSimuFuncSize", &StOpt::OptimizerDPBase::getSimuFuncSize)
    ;

    // map optimizer DP base for tree
    //******************************
    py::class_<StOpt::OptimizerDPTreeBase, std::shared_ptr<StOpt::OptimizerDPTreeBase>, PyOptimizerDPTreeBase>(m, "OptimizerDPTreeBase")
    .def("getCone",  [](const StOpt::OptimizerDPTreeBase & p_x, const  std::vector< Eigen::ArrayXd >    &p_region)
    {
        std::vector< std::array< double, 2> > region;
        region.reserve(p_region.size());
        for (const auto &item : p_region)
        {
            region.push_back(std::array< double, 2> {{ item[0], item[1]}});
        }
        std::vector< std::array< double, 2> > resLoc = p_x.getCone(region);
        std::vector< Eigen::ArrayXd > result;
        result.reserve(resLoc.size());
        for (const auto &item : resLoc)
        {
            Eigen::ArrayXd dom(2);
            dom(0) = item[0];
            dom(1) = item[1];
            result.push_back(dom);
        }
        return result;
    }
        )
    .def("getDimensionToSplit",  &StOpt::OptimizerDPTreeBase::getDimensionToSplit)
    .def("stepOptimize", [](const StOpt::OptimizerDPTreeBase & p_x,
                            const std::shared_ptr< StOpt::SpaceGrid> &p_grid,
                            const Eigen::ArrayXd   & p_reg,
                            const std::vector<StOpt::ContinuationValueTree> &p_cont)
    {
        std::pair< Eigen::ArrayXXd, Eigen::ArrayXXd>  resLoc = p_x.stepOptimize(p_grid, p_reg, p_cont);
        py::list res ;
        res.append(py::cast(resLoc.first));
        res.append(py::cast(resLoc.second));
        return res;
    }
        )
    .def("stepSimulate", &StOpt::OptimizerDPTreeBase::stepSimulate)
    .def("stepSimulateControl", &StOpt::OptimizerDPTreeBase::stepSimulateControl)
    .def("getNbRegime", &StOpt::OptimizerDPTreeBase::getNbRegime)
    .def("getNbControl", &StOpt::OptimizerDPTreeBase::getNbControl)
    .def("getSimulator", &StOpt::OptimizerDPTreeBase::getSimulator)
    .def("getSimuFuncSize", &StOpt::OptimizerDPTreeBase::getSimuFuncSize)
    ;


    // map optimizer switching base
    //**********************
    py::class_<StOpt::OptimizerSwitchBase, std::shared_ptr<StOpt::OptimizerSwitchBase>, PyOptimizerSwitchBase>(m, "OptimizerSwitchBase")
    .def("getCone", &StOpt::OptimizerSwitchBase::getCone)
    .def("getDimensionToSplit",  &StOpt::OptimizerSwitchBase::getDimensionToSplit)
    .def("stepOptimize", [](const StOpt::OptimizerSwitchBase & p_x,
                            std::vector<std::shared_ptr< StOpt::RegularSpaceIntGrid> >   &p_grids,
                            const   int &p_iReg,
                            const   Eigen::ArrayXi  & p_state,
                            const   std::shared_ptr< StOpt::BaseRegression>   &p_condExp,
                            const std::vector < Eigen::ArrayXXd  > &p_val)
    {
        std::vector < std::shared_ptr< Eigen::ArrayXXd > > val;
        val.reserve(p_val.size());
        for (const auto &item : p_val)
            val.push_back(std::make_shared<Eigen::ArrayXXd>(item));
        Eigen::ArrayXXd  resLoc = p_x.stepOptimize(p_grids, p_iReg, p_state, p_condExp, val);
        return resLoc;
    }
        )
    .def("getNbRegime", &StOpt::OptimizerSwitchBase::getNbRegime)
    .def("getSimulator", &StOpt::OptimizerSwitchBase::getSimulator)
    .def("getSimuFuncSize", &StOpt::OptimizerSwitchBase::getSimuFuncSize)
    ;

    // map final pay off
    //******************
    py::class_<StOpt::PayOffBase, std::shared_ptr<StOpt::PayOffBase>  >(m, "PayOffBase");

    // // maps final step
    py::class_<PyFinalStepDP, std::shared_ptr<PyFinalStepDP>  >(m, "FinalStepDP")
    .def(py::init< const std::shared_ptr<StOpt::SpaceGrid>  &, const int & >())
    .def("set", &PyFinalStepDP::set)
    .def("_setItem_", &PyFinalStepDP::operator())
    ;


    /// map one time step transition
    py::class_<PyTransitionStepRegressionDP, std::shared_ptr<PyTransitionStepRegressionDP> >(m, "TransitionStepRegressionDP")
    .def(py::init< const std::shared_ptr<StOpt::SpaceGrid> &,
         const  std::shared_ptr<StOpt::SpaceGrid> &, const std::shared_ptr<StOpt::OptimizerDPBase > &>())
    .def("oneStep", & PyTransitionStepRegressionDP::oneStepWrap)
    .def("dumpContinuationValues", & PyTransitionStepRegressionDP::dumpContinuationValuesWrap)
    .def("dumpBellmanValues", & PyTransitionStepRegressionDP::dumpBellmanValuesWrap)
    ;

    // maps SimulateStepRegression for simulation purpose
    py::class_<PySimulateStepRegression, std::shared_ptr<PySimulateStepRegression> >(m, "SimulateStepRegression")
    .def(py::init< BinaryFileArchiveStOpt &,  const int &,  const std::string &,
         const   std::shared_ptr< StOpt::SpaceGrid > &, const  std::shared_ptr< StOpt::OptimizerDPBase > & >())
    .def("oneStep", &PySimulateStepRegression::oneStepWrap)
    ;


    py::class_<PySimulateStepRegressionControl, std::shared_ptr<PySimulateStepRegressionControl> >(m, "SimulateStepRegressionControl")
    .def(py::init< BinaryFileArchiveStOpt &,  const int &,  const std::string &,
         const   std::shared_ptr< StOpt::SpaceGrid > &, const  std::shared_ptr< StOpt::OptimizerDPBase > & >())
    .def("oneStep", & PySimulateStepRegressionControl::oneStepWrap)
    ;


    /// map one time step transition by tree
    py::class_<PyTransitionStepTreeDP, std::shared_ptr<PyTransitionStepTreeDP> >(m, "TransitionStepTreeDP")
    .def(py::init< const std::shared_ptr<StOpt::SpaceGrid> &,
         const  std::shared_ptr<StOpt::SpaceGrid> &, const std::shared_ptr<StOpt::OptimizerDPTreeBase > &>())
    .def("oneStep", & PyTransitionStepTreeDP::oneStepWrap)
    .def("dumpContinuationValues", & PyTransitionStepTreeDP::dumpContinuationValuesWrap)
    ;

    // maps SimulateStepTree for simulation purpose, by tree
    py::class_<PySimulateStepTree, std::shared_ptr<PySimulateStepTree> >(m, "SimulateStepTree")
    .def(py::init< BinaryFileArchiveStOpt &,  const int &,  const std::string &,
         const   std::shared_ptr< StOpt::SpaceGrid > &, const  std::shared_ptr< StOpt::OptimizerDPTreeBase > & >())
    .def("oneStep", &PySimulateStepTree::oneStepWrap)
    ;

    // maps simulation using control for trees
    py::class_<PySimulateStepTreeControl, std::shared_ptr<PySimulateStepTreeControl> >(m, "SimulateStepTreeControl")
    .def(py::init< BinaryFileArchiveStOpt &,  const int &,  const std::string &,
         const   std::shared_ptr< StOpt::SpaceGrid > &, const  std::shared_ptr< StOpt::OptimizerDPTreeBase > & >())
    .def("oneStep", & PySimulateStepTreeControl::oneStepWrap)
    ;

    // maping for switching
    py::class_<PyTransitionStepRegressionSwitch, std::shared_ptr< PyTransitionStepRegressionSwitch> >(m, "TransitionStepRegressionSwitch")
    .def(py::init< const  std::vector< std::shared_ptr<StOpt::RegularSpaceIntGrid> >  &,  const  std::vector< std::shared_ptr<StOpt::RegularSpaceIntGrid> >  &,
         const  std::shared_ptr<StOpt::OptimizerSwitchBase > & >())
    .def("oneStep", &PyTransitionStepRegressionSwitch::oneStepWrap)
    .def("dumpContinuationValues", &PyTransitionStepRegressionSwitch::dumpContinuationValuesWrap)
    ;

    // maps SimulateStepRegression for simulation purpose
    py::class_<PySimulateStepSwitch, std::shared_ptr<PySimulateStepSwitch> >(m, "SimulateStepSwitch")
    .def(py::init< BinaryFileArchiveStOpt &,  const int &,  const std::string &,
         const   std::vector<std::shared_ptr< StOpt::RegularSpaceIntGrid > > &, const  std::shared_ptr< StOpt::OptimizerSwitchBase > & >())
    .def("oneStep", &PySimulateStepSwitch::oneStepWrap)
    ;


#ifdef USE_MPI

    py::class_<PyFinalStepDPDist, std::shared_ptr<PyFinalStepDPDist>  >(m, "FinalStepDPDist")
    .def(py::init< const std::shared_ptr<StOpt::SpaceGrid>  &, const int &, const Eigen::Array< bool, Eigen::Dynamic, 1>  & >())
    .def("set", &PyFinalStepDPDist::set)
    .def("_setItem_", &PyFinalStepDPDist::operator())
    ;

    py::class_<PyFinalStepZeroDist, std::shared_ptr<PyFinalStepZeroDist>  >(m, "FinalStepZeroDist")
    .def(py::init<const  std::vector< std::shared_ptr<StOpt::RegularSpaceIntGrid > >   &,
         const  std::vector< Eigen::Array< bool, Eigen::Dynamic, 1> >   &  >())
    .def("_setItem_", &PyFinalStepZeroDist::operator())
    ;

    py::class_<PyTransitionStepRegressionDPDist, std::shared_ptr<PyTransitionStepRegressionDPDist> >(m, "TransitionStepRegressionDPDist")
    .def(py::init< const std::shared_ptr<StOpt::SpaceGrid> &,
         const  std::shared_ptr<StOpt::SpaceGrid> &, const std::shared_ptr<StOpt::OptimizerDPBase > &>())
    .def("oneStep", & PyTransitionStepRegressionDPDist::oneStepWrap)
    .def("dumpContinuationValues", & PyTransitionStepRegressionDPDist::dumpContinuationValuesWrap)
    .def("dumpBellmanValues", & PyTransitionStepRegressionDPDist::dumpBellmanValuesWrap)
    ;

    py::class_<PySimulateStepRegressionDist, std::shared_ptr<PySimulateStepRegressionDist> >(m, "SimulateStepRegressionDist")
    .def(py::init< BinaryFileArchiveStOpt &,  const int &,  const std::string &,
         const   std::shared_ptr< StOpt::SpaceGrid > &, const  std::shared_ptr< StOpt::OptimizerDPBase > &, const bool & >())
    .def("oneStep", &PySimulateStepRegressionDist::oneStepWrap)
    ;


    py::class_<PySimulateStepRegressionControlDist, std::shared_ptr<PySimulateStepRegressionControlDist> >(m, "SimulateStepRegressionControlDist")
    .def(py::init< BinaryFileArchiveStOpt &,  const int &,  const std::string &, const   std::shared_ptr< StOpt::SpaceGrid > &,
         const   std::shared_ptr< StOpt::SpaceGrid > &, const  std::shared_ptr< StOpt::OptimizerDPBase > &, const bool & >())
    .def("oneStep", & PySimulateStepRegressionControlDist::oneStepWrap)
    ;

    py::class_<PyTransitionStepTreeDPDist, std::shared_ptr<PyTransitionStepTreeDPDist> >(m, "TransitionStepTreeDPDist")
    .def(py::init< const std::shared_ptr<StOpt::SpaceGrid> &,
         const  std::shared_ptr<StOpt::SpaceGrid> &, const std::shared_ptr<StOpt::OptimizerDPTreeBase > &>())
    .def("oneStep", & PyTransitionStepTreeDPDist::oneStepWrap)
    .def("dumpContinuationValues", & PyTransitionStepTreeDPDist::dumpContinuationValuesWrap)
    ;

    py::class_<PySimulateStepTreeDist, std::shared_ptr<PySimulateStepTreeDist> >(m, "SimulateStepTreeDist")
    .def(py::init< BinaryFileArchiveStOpt &,  const int &,  const std::string &,
         const   std::shared_ptr< StOpt::SpaceGrid > &, const  std::shared_ptr< StOpt::OptimizerDPTreeBase > &, const bool & >())
    .def("oneStep", &PySimulateStepTreeDist::oneStepWrap)
    ;


    py::class_<PySimulateStepTreeControlDist, std::shared_ptr<PySimulateStepTreeControlDist> >(m, "SimulateStepTreeControlDist")
    .def(py::init< BinaryFileArchiveStOpt &,  const int &,  const std::string &, const   std::shared_ptr< StOpt::SpaceGrid > &,
         const   std::shared_ptr< StOpt::SpaceGrid > &, const  std::shared_ptr< StOpt::OptimizerDPTreeBase > &, const bool & >())
    .def("oneStep", & PySimulateStepTreeControlDist::oneStepWrap)
    ;

    // maping for switching
    py::class_<PyTransitionStepRegressionSwitchDist, std::shared_ptr< PyTransitionStepRegressionSwitchDist> >(m, "TransitionStepRegressionSwitchDist")
    .def(py::init< const  std::vector< std::shared_ptr<StOpt::RegularSpaceIntGrid> >  &,  const  std::vector< std::shared_ptr<StOpt::RegularSpaceIntGrid> >  &,
         const  std::shared_ptr<StOpt::OptimizerSwitchBase > & >())
    .def("oneStep", &PyTransitionStepRegressionSwitchDist::oneStepWrap)
    .def("dumpContinuationValues", &PyTransitionStepRegressionSwitchDist::dumpContinuationValuesWrap)
    ;

    m.def("reconstructProc0Mpi", & pyReconstructProc0Mpi);
    m.def("reconstructProc0ForIntMpi", & pyReconstructProc0ForIntMpi);

#endif

}
