# Copyright (C) 2016 EDF
# All Rights Reserved
# This code is published under the GNU Lesser General Public License (GNU LGPL)
import numpy as np
import math
import StOptReg as reg
import StOptGrids
import StOptGlobal
import Simulators as sim
import Optimizers as opt
import Utils
import dp.DynamicProgrammingByRegressionHighLevel as dyn
import dp.SimulateRegressionControlHighLevel as srt
import unittest
import importlib

accuracyClose = 0.5


def testGasStorage():

    # test MPI
    moduleMpi4Py=importlib.util.find_spec('mpi4py')
    if (moduleMpi4Py is not None):
        from mpi4py import MPI
        # storage
    ###############
    maxLevelStorage = 360000
    injectionRateStorage = 60000
    withdrawalRateStorage = 45000
    injectionCostStorage = 0.35
    withdrawalCostStorage = 0.35
    switchingCostStorage =  4.

    maturity = 1.
    nstep = 10

    # define a a time grid
    timeGrid = StOptGrids.OneDimRegularSpaceGrid(0., maturity / nstep, nstep)
    futValues = []
    # periodicity factor
    iPeriod = 52

    for i in 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
    nDim = 1
    sigma = np.zeros(nDim) + 0.94
    mr = np.zeros(nDim) + 0.29
    # number of simulations
    nbsimulOpt = 20000
    # grid
    #####
    nGrid = 40
    lowValues = np.zeros(1, dtype = float)
    step = np.zeros(1, dtype = float) + maxLevelStorage / nGrid
    nbStep = np.zeros(1, dtype = np.int32) + nGrid
    grid = StOptGrids.RegularSpaceGrid(lowValues, step, nbStep)

    # no actualization
    rate=0.
    # a backward simulator
    ######################
    bForward = False
    backSimulator = sim.MeanRevertingSimulator(futureGrid, sigma, mr, rate,maturity, nstep, nbsimulOpt, bForward)
    # optimizer
    ############
    storage = opt.OptimizeGasStorageSwitchingCostMeanReverting(injectionRateStorage, withdrawalRateStorage, injectionCostStorage, withdrawalCostStorage, switchingCostStorage)
    # regressor
    ##########
    nMesh = 4
    nbMesh = np.zeros(1, dtype = np.int32) + nMesh
    regressor = reg.LocalLinearRegression(nbMesh)
    # final value
    vFunction = Utils.ZeroPayOff()

    # initial values
    initialStock = np.zeros(1) + maxLevelStorage
    initialRegime = 0 # here do nothing (no injection, no withdrawal)

    # Optimize
    ###########
    fileToDump = "CondExpGasSwiCostHL"
    # link the simulations to the optimizer
    storage.setSimulator(backSimulator)
    valueOptim = dyn.DynamicProgrammingByRegressionHighLevel(grid, storage, regressor, vFunction, initialStock, initialRegime, fileToDump)
    print("valOP", valueOptim)
    nbsimulSim = 40000
    bForward = True
    forSimulator2 = sim.MeanRevertingSimulator(futureGrid, sigma, mr, rate,maturity, nstep, nbsimulSim, bForward)
    storage.setSimulator(forSimulator2)
    valSimu2 = srt.SimulateRegressionControl(grid, storage, vFunction, initialStock, initialRegime, fileToDump)
    print("valSimu2", valSimu2, "valOptim", valueOptim)

class testGasStorageSwitchCostTest(unittest.TestCase):

        def test_gasStorageSwitchCostHighLevel(self):

            testGasStorage()

if __name__ == '__main__':

        unittest.main()











