# Copyright (C) 2021 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.geners as StOptGeners

def DynamicProgrammingSwitchingByRegressionHighLevel(p_grid, p_optimize, p_regressor, p_pointState, p_initialRegime, p_fileToDump) :
    
    # from the optimizer get back the simulation
    simulator = p_optimize.getSimulator()
    print("simulator", simulator)
    # final values
    valuesNext= []
    for iReg in range(len(p_grid)):
        valuesNext.append( np.zeros([simulator.getNbSimul(),p_grid[iReg].getNbPoints()]))
        
    ar = StOptGeners.BinaryFileArchive(p_fileToDump, "w")
    nameAr = "ContinuationSwitching"
    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 = StOptGlobal.TransitionStepRegressionSwitch(p_grid, p_grid, p_optimize)
        values = transStep.oneStep(valuesNext, p_regressor)
        transStep.dumpContinuationValues(ar, nameAr, nsteps - 1 -iStep, valuesNext, p_regressor)
        valuesNext = values
        
    # interpolate at the initial stock point and initial regime
    return valuesNext[p_initialRegime][:,p_grid[p_initialRegime].globCoordPerDimToLocal( p_pointState)].mean()
