# 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


# Simulate the optimal strategy , threaded version, for tree
# p_grid                   grid used for  deterministic state (stocks for example)
# p_optimize               optimizer defining the optimization between two time steps
# p_funcFinalValue         function defining the final value
# p_pointStock             initial point stock
# p_initialRegime          regime at initial date
# p_fileToDump             name of the file used to dump continuation values in optimization
def SimulateTreeHighLevel(p_grid, p_optimize, p_funcFinalValue, p_pointStock, p_initialRegime, p_fileToDump) :

    simulator = p_optimize.getSimulator()
    nbStep = simulator.getNbStep()
    states = []
    for i in range(simulator.getNbSimul()) :
        states.append(StOptGlobal.StateTreeStocks(p_initialRegime, p_pointStock, 0))
    ar = StOptGeners.BinaryFileArchive(p_fileToDump, "r")
    # name for continuation object in archive
    nameAr = "Continuation"
    # cost function
    costFunction = np.zeros((p_optimize.getSimuFuncSize(), simulator.getNbSimul()))
    
    # iterate on time steps
    for istep in range(nbStep) :
        NewState = StOptGlobal.SimulateStepTree(ar, nbStep - 1 - istep, nameAr, p_grid, p_optimize).oneStep(states, costFunction)
        # different from C++
        states = NewState[0]
        costFunction = NewState[1]
     
        # new stochastic state
        simulator.stepForward()
        
        for i in range(simulator.getNbSimul()) :
            states[i].setStochasticRealization(simulator.getNodeAssociatedToSim(i))
    # final : accept to exercise if not already done entirely
    for i in range(simulator.getNbSimul()) :
        costFunction[0,i] += p_funcFinalValue.set(states[i].getRegime(), states[i].getPtStock(),simulator.getValueAssociatedToNode(states[i].getStochasticRealization()))
        
    # average gain/cost
    return costFunction.mean()
