# Copyright (C) 2019 EDF
# All Rights Reserved
# This code is published under the GNU Lesser General Public License (GNU LGPL)
import numpy as np
import pystopt.grids as StOptGrids
import pystopt.geners as StOptGeners
import pystopt.glob as StOptGlobal
import pystopt.tree as StOptTree


def DynamicProgrammingByTreeHighLevel(p_grid, p_optimize,  p_funcFinalValue, p_pointStock, p_initialRegime, p_fileToDump) :
    
    # from the optimizer get back the simulation
    simulator = p_optimize.getSimulator()
    # final values
    fin = StOptGlobal.FinalStepDP(p_grid, p_optimize.getNbRegime())
    valuesNext = fin.set(p_funcFinalValue, simulator.getNodes())
    ar = StOptGeners.BinaryFileArchive(p_fileToDump, "w")
    nameAr = "Continuation"
    nsteps =simulator.getNbStep()
    # iterate on time steps
    for iStep in range(nsteps) :
        simulator.stepBackward()
        # get back probability
        proba = simulator.getProba()
        # and connection matrix
        connected = simulator.getConnected()
        # creta tree for conditional expectation
        tree=StOptTree.Tree(proba,connected)      
        # transition object
        transStep = StOptGlobal.TransitionStepTreeDP(p_grid, p_grid, p_optimize)
        valuesAndControl = transStep.oneStep(valuesNext, tree)
        transStep.dumpContinuationValues(ar, nameAr, iStep, valuesNext, valuesAndControl[1], tree)
        valuesNext = valuesAndControl[0]
        
    # interpolate at the initial stock point and initial regime
    return (p_grid.createInterpolator(p_pointStock).applyVec(valuesNext[p_initialRegime])).mean()
