# 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 math
import StOptGrids
import StOptGlobal
import StOptGeners
import Simulators
import Optimizers
import Utils
import dp.DynamicProgrammingByTreeHighLevel as dyn
import dp.SimulateTreeControlHighLevel as simC
import dp.SimulateTreeHighLevel as simV
import unittest
import importlib

accuracyClose = 1.5


class ZeroFunction :
    
    def __init__(self) :
        
        return None

    def set(self, a, b, c) :
        
        return 0.

# valorization of a given gas storage on a  grid
# p_grid             the grid
# p_maxLevelStorage  maximum level
def gasStorage(p_grid, p_maxLevelStorage) :
    
    # storage
    injectionRateStorage = 60000
    withdrawalRateStorage = 45000
    injectionCostStorage = 0.35
    withdrawalCostStorage = 0.35

    maturity = 1.
    nstep = 100

    # define a a time grid
    timeGrid = StOptGrids.OneDimRegularSpaceGrid(0., maturity / nstep, nstep)
    # future values
    futValues = []

    # periodicity factor
    iPeriod = 52

    for i in list(range(nstep + 1)) :
        futValues.append(50. + 20. * math.sin((math.pi * i * iPeriod) / nstep))
      
    # define the future curve
    futureGrid = Utils.FutureCurve(timeGrid, futValues)
    
    # one dimensional factors
    mr = 0.29
    sigma =0.94

    # step
    dt = 1. / 100.
        
    # simulation dates
    dates =  dt * np.arange(0,nstep+1)

    # simulation dates
    tree = Simulators.TrinomialTreeOUSimulator(mr, sigma, dates)

    indexT = np.arange(0,nstep+1)
    tree.dump("Tree",indexT)

    # read archive
    archiveToRead = StOptGeners.BinaryFileArchive("Tree", "r");

    # backward simulator
    backSimulator = Simulators.MeanRevertingSimulatorTree(archiveToRead, futureGrid, sigma, mr)


    # optimizer
    storage = Optimizers.OptimizeGasStorageMeanRevertingTree(injectionRateStorage, withdrawalRateStorage, injectionCostStorage, withdrawalCostStorage)
    
    # final value
    vFunction = Utils.ZeroPayOff()
    
    # initial values
    initialStock = np.zeros(1) + p_maxLevelStorage
    initialRegime = 0 # only one regime
    
    # Optimize
    fileToDump = "CondExpGasStorage"
    
    # link the simulations to the optimizer
    storage.setSimulator(backSimulator)
    valueOptim = dyn.DynamicProgrammingByTreeHighLevel(p_grid, storage,  vFunction, initialStock, initialRegime, fileToDump)
    print("valueOptim", valueOptim)

    # new simulator for simulate optimal control
    nbsimulSim = 10000
    forwardSimulator = Simulators.MeanRevertingSimulatorTree(archiveToRead, futureGrid, sigma, mr,nbsimulSim)
     # link the simulations to the optimizer
    storage.setSimulator(forwardSimulator)

    # forward simulation of optimal control
    valueSim =simV.SimulateTreeHighLevel(p_grid,storage,vFunction, initialStock, initialRegime, fileToDump)
    print("Value simulation ", valueSim)

    # same with control
    forwardSimulator = Simulators.MeanRevertingSimulatorTree(archiveToRead, futureGrid, sigma, mr,nbsimulSim)
     # link the simulations to the optimizer
    storage.setSimulator(forwardSimulator)

    # forward simulation of optimal control
    valueSim =simC.SimulateTreeControlHighLevel(p_grid,storage,vFunction, initialStock, initialRegime, fileToDump)
    print("Value simulation with control", valueSim)
    
    return valueOptim



class testGasStorageTest(unittest.TestCase):
          
    def test_simpleStorage(self):
              
        # test MPI
        moduleMpi4Py=importlib.util.find_spec('mpi4py')
        if (moduleMpi4Py is not None):
            from mpi4py import MPI
        # storage
        #########
        maxLevelStorage = 90000
        # grid
        ######
        nGrid = 10
        lowValues = np.zeros(1)
        step = np.zeros(1) + (maxLevelStorage / nGrid)
        nbStep = np.zeros(1, dtype = np.int32) + nGrid
        grid = StOptGrids.RegularSpaceGrid(lowValues, step, nbStep)
              
        gasStorage(grid, maxLevelStorage)
     

        
if __name__ == '__main__':
    unittest.main()
