# Copyright (C) 2016 EDF
# All Rights Reserved
# This code is published under the GNU Lesser General Public License (GNU LGPL)
import StOptGrids 
import StOptReg
import StOptGlobal
import StOptGeners
import imp
import numpy as np


def DynamicProgrammingByRegressionDist(p_grid, p_optimize, p_regressor, p_funcFinalValue, p_pointStock, p_initialRegime, p_fileToDump, p_bOneFile) :
    
    try:
        imp.find_module('mpi4py')
        found = True
    except:
        print("Not parallel module found ")
        found = False
        
    if found :
        # 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()
