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


def DynamicProgrammingByRegressionDist(p_grid, p_optimize, p_regressor, p_funcFinalValue, p_pointStock, p_initialRegime, p_fileToDump, p_bOneFile) :
    
    moduleMpi4Py=importlib.util.find_spec('mpi4py')
    if (moduleMpi4Py is not None):
        from mpi4py import MPI
        world = MPI.COMM_WORLD
        # from the optimizer get back the simulation
        simulator = p_optimize.getSimulator()
        fin = StOptGlobal.FinalStepDPDist(p_grid, p_optimize.getNbRegime(),p_optimize.getDimensionToSplit())
        # final values
        valuesNext = fin.set(p_funcFinalValue, np.asarray(simulator.getParticles()))
        
        from mpi4py import MPI
        world = MPI.COMM_WORLD
        # dump
        toDump = p_fileToDump
        
        # test if one file generated
        if p_bOneFile is False:
            toDump += "_" + str(world.rank)
            
        ar = StOptGeners.BinaryFileArchive(toDump, "w")
        # name for object in archive
        nameAr = "Continuation"
        
        nsteps =simulator.getNbStep()
        # iterate on time steps
        for iStep in range(nsteps) :
            asset = simulator.stepBackwardAndGetParticles()
            
            # conditional espectation operator
            if iStep == (simulator.getNbStep() - 1):
                p_regressor.updateSimulations(True, asset)
            
            else:
                p_regressor.updateSimulations(False, asset)
            
            # transition object
            transStep = StOptGlobal.TransitionStepRegressionDPDist(p_grid, p_grid, p_optimize)
            valuesAndControl = transStep.oneStep(valuesNext, p_regressor)
            transStep.dumpContinuationValues(ar, nameAr,  nsteps - 1 - iStep, valuesNext, valuesAndControl[1], p_regressor, p_bOneFile)
            valuesNext = valuesAndControl[0]
            
        # reconstruct a small grid for interpolation
        return StOptGlobal.reconstructProc0Mpi(p_pointStock, p_grid, valuesNext[p_initialRegime],p_optimize.getDimensionToSplit()).mean()
