# Copyright (C) 2016, 2018 EDF
# All Rights Reserved
# This code is published under the GNU Lesser General Public License (GNU LGPL)
import pystopt.regression as reg
import pystopt.geners as StOptGeners
import dp.TransitionStepRegressionDP as trans
import dp.FinalStepDP as final


def DynamicProgrammingByRegression(p_grid, p_optimize, p_regressor, p_funcFinalValue, p_pointStock, p_initialRegime, p_fileToDump, key1="Continuation" , key2 = "Control"):

    # from the optimizer get back the simulation
    simulator = p_optimize.getSimulator()
    # final values
    valuesNext = final.FinalStepDP(p_grid, p_optimize.getNbRegime()).operator(p_funcFinalValue, simulator.getParticles())

    archiveToWrite = StOptGeners.BinaryFileArchive(p_fileToDump, "w")
    nsteps = simulator.getNbStep()
    # iterate on time steps
    for iStep in range(nsteps):
        asset = simulator.stepBackwardAndGetParticles()

        # conditional expectation operator
        if iStep == (simulator.getNbStep() - 1):
            p_regressor.updateSimulations(True, asset)
        else:
            p_regressor.updateSimulations(False, asset)

        # transition object
        transStep = trans.TransitionStepRegressionDP(p_grid, p_grid, p_optimize)
        valuesAndControl = transStep.oneStep(valuesNext, p_regressor)
        # Dump the continuation values in the archive:
        control =  valuesAndControl[1]
        archiveToWrite.dumpGridAndRegressedValue(key1, nsteps - 1 - iStep, valuesNext, p_regressor, p_grid)
        archiveToWrite.dumpGridAndRegressedValue(key2, nsteps - 1 - iStep, control, p_regressor, p_grid)
        valuesNext = valuesAndControl[0]


    # interpolate at the initial stock point and initial regime
    return (p_grid.createInterpolator(p_pointStock).applyVec(valuesNext[p_initialRegime])).mean()
