#include <iostream>
#include "StOpt/core/utils/constant.h"
#include "test/c++/tools/semilagrangien/OptimizeSLCase3.h"

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

OptimizerSLCase3::OptimizerSLCase3(const double &p_mu, const double &p_sig, const double &p_dt, const double &p_alphaMax, const double &p_stepAlpha):
    m_dt(p_dt), m_mu(p_mu), m_sig(p_sig), m_alphaMax(p_alphaMax), m_stepAlpha(p_stepAlpha) {}

vector< array< double, 2> >  OptimizerSLCase3::getCone(const  vector<  array< double, 2>  > &p_xInit) const
{
    vector< array< double, 2> > xReached(1);
    xReached[0][0] = p_xInit[0][0] -  m_alphaMax * m_mu / m_sig * m_dt -  m_alphaMax * sqrt(m_dt);
    xReached[0][1] = p_xInit[0][1]  + m_alphaMax * sqrt(m_dt) ;
    return xReached;
}

pair< ArrayXd, ArrayXd> OptimizerSLCase3::stepOptimize(const ArrayXd   &p_point,
        const vector< shared_ptr<SemiLagrangEspCond> > &p_semiLag, const double &, const Eigen::ArrayXd &) const
{
    pair< ArrayXd, ArrayXd> solutionAndControl;
    solutionAndControl.first.resize(1);
    solutionAndControl.second.resize(1);
    ArrayXd b(1);
    ArrayXXd sig(1, 1) ;
    double vMin = StOpt::infty;
    for (int iAl = 0; iAl < m_alphaMax / m_stepAlpha; ++iAl)
    {
        double alpha = iAl * m_stepAlpha;
        b(0) = -alpha * m_mu / m_sig; // trend
        sig(0) = alpha; // volatility with one Brownian
        pair<double, bool> lagrang = p_semiLag[0]->oneStep(p_point, b, sig, m_dt); // test the control
        if (lagrang.second)
        {
            if (lagrang.first < vMin)
            {
                vMin = lagrang.first;
                solutionAndControl.second(0) = alpha;
            }
        }
    }

    solutionAndControl.first(0) =  vMin;
    return solutionAndControl;
}

void OptimizerSLCase3::stepSimulate(const StOpt::SpaceGrid   &p_gridNext,
                                    const std::vector< shared_ptr< StOpt::SemiLagrangEspCond > > &p_semiLag,
                                    Eigen::Ref<Eigen::ArrayXd>  p_state,  int &,
                                    const Eigen::ArrayXd &p_gaussian, const Eigen::ArrayXd &,
                                    Eigen::Ref<Eigen::ArrayXd>) const
{
    double vMin = StOpt::infty;
    double alphaOpt = -1;
    ArrayXd b(1);
    ArrayXXd sig(1, 1) ;
    ArrayXd proba = p_state ;
    // recalculate the optimal alpha
    for (int iAl = 0; iAl < m_alphaMax / m_stepAlpha; ++iAl)
    {
        double alpha = iAl * m_stepAlpha;
        b(0) = -alpha * m_mu / m_sig;// trend
        sig(0) = alpha;// volatility with one Brownian
        pair<double, bool> lagrang = p_semiLag[0]->oneStep(proba, b, sig, m_dt);// test the control
        if (lagrang.second)
        {
            if (lagrang.first < vMin)
            {
                vMin = lagrang.first;
                alphaOpt = alpha;
            }
        }
    }
    proba(0) +=  alphaOpt * p_gaussian(0) * sqrt(m_dt);
    // truncate if necessary
    p_gridNext.truncatePoint(proba);
    p_state = proba ;

}


void OptimizerSLCase3:: stepSimulateControl(const SpaceGrid    &p_gridNext,
        const  vector< shared_ptr< InterpolatorSpectral> >   &p_controlInterp,
        Eigen::Ref<Eigen::ArrayXd>  p_state,   int &,
        const ArrayXd &p_gaussian,
        Eigen::Ref<Eigen::ArrayXd>) const
{
    ArrayXd proba = p_state ;
    double alphaOpt = p_controlInterp[0]->apply(p_state);
    proba(0) +=  alphaOpt * p_gaussian(0) * sqrt(m_dt);
    // truncate if necessary
    p_gridNext.truncatePoint(proba);
    p_state = proba ;
}
