# 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 StOptGrids 
import StOptReg
import StOptGlobal
import 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()
