# Copyright (C) 2016, 2018 EDF
# All Rights Reserved
# This code is published under the GNU Lesser General Public License (GNU LGPL)
import StOptReg
import 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()
