#!/usr/bin/python3

# 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 Utils
import Simulators as sim
import Optimizers as opt
import dp.DynamicProgrammingByRegressionDist as dynmpi
import dp.SimulateRegressionControlDist as srtmpi
import unittest
import importlib

accuracyClose = 1e6
accuracyEqual = 0.0001

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

    maturity = 1.
    nstep = 10

    # 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
    nDim = 1
    sigma = np.zeros(nDim) + 0.94
    mr = np.zeros(nDim) + 0.29
    # number of simulations
    nbsimulOpt = 20000

    # no actualization
    rate =0.
    # a backward simulator
    bForward = False

    backSimulator = sim.MeanRevertingSimulator(futureGrid, sigma, mr, rate, maturity, nstep, nbsimulOpt, bForward)
    # optimizer                  
    storage = opt.OptimizeGasStorageMeanReverting(injectionRateStorage, withdrawalRateStorage, injectionCostStorage, withdrawalCostStorage)

    # regressor
    nMesh = 6
    nbMesh = np.zeros(1, dtype = np.int32) + nMesh
    regressor = reg.LocalLinearRegression(nbMesh)
    # final value
    vFunction = Utils.ZeroPayOff()

    # initial values
    initialStock = np.zeros(1) + p_maxLevelStorage
    initialRegime = 0 # only one regime

    # Optimize
    fileToDump = "CondExpGasStorageHLMpi"
    bOneFile = True

    # link the simulations to the optimizer
    storage.setSimulator(backSimulator)
    valueOptimMpi = dynmpi.DynamicProgrammingByRegressionDist(p_grid, storage, regressor, vFunction, initialStock, initialRegime, fileToDump, bOneFile)
    print("valOptimMpi", valueOptimMpi)

    world.barrier()

    nbsimulSim = 40000
    bForward = True
    forSimulator = sim.MeanRevertingSimulator(futureGrid, sigma, mr, rate,maturity, nstep, nbsimulSim, bForward)
    storage.setSimulator(forSimulator)
    valSimuMpi = srtmpi.SimulateRegressionControlDist(p_grid, storage, vFunction, initialStock, initialRegime, fileToDump, bOneFile)
    print("valSimuMpi", valSimuMpi)

    return valueOptimMpi, valSimuMpi
                    
class testGasStorageTest(unittest.TestCase):
           
    def test_simpleStorageMpi(self):
               
        moduleMpi4Py=importlib.util.find_spec('mpi4py')
        if (moduleMpi4Py is not None):
            from mpi4py import MPI
            world = MPI.COMM_WORLD
            # 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)
                   
            val = gasStorage(grid, maxLevelStorage)
            
            if world.rank == 0:
                self.assertAlmostEqual(val[0], val[1], None, None, accuracyClose)    
            
            world.barrier()
            
            # grid
            ######
            poly = np.zeros(1, dtype = np.int32) + 1
            gridL = StOptGrids.RegularLegendreGrid(lowValues, step, nbStep, poly)
            
            valLegendre = gasStorage(gridL, maxLevelStorage)
            
            if world.rank == 0:
                self.assertAlmostEqual(valLegendre[0], valLegendre[1], None, None, accuracyClose)
         
    def test_simpleStorageLegendreQuadratic(self):
         
        moduleMpi4Py=importlib.util.find_spec('mpi4py')
        if (moduleMpi4Py is not None):
            from mpi4py import MPI
            world = MPI.COMM_WORLD
            # storage
            maxLevelStorage = 90000
            # grid
            nGrid = 5
            lowValues = np.zeros(1)
            step = np.zeros(1) + (maxLevelStorage / nGrid)
            nbStep = np.zeros(1, dtype = np.int32) + nGrid
            poly = np.zeros(1, dtype = np.int32) + 2
            grid = StOptGrids.RegularLegendreGrid(lowValues, step, nbStep, poly)
             
            val = gasStorage(grid, maxLevelStorage)
            
            if world.rank == 0:
                self.assertAlmostEqual(val[0], val[1], None, None, accuracyClose)
         
    def test_simpleStorageLegendreCubic(self):
         
        moduleMpi4Py=importlib.util.find_spec('mpi4py')
        if (moduleMpi4Py is not None):
            from mpi4py import MPI
            world = MPI.COMM_WORLD
            # storage
            maxLevelStorage = 90000
            # grid
            nGrid = 5
            lowValues = np.zeros(1)
            step = np.zeros(1) + (maxLevelStorage / nGrid)
            nbStep = np.zeros(1, dtype = np.int32) + nGrid
            poly = np.zeros(1, dtype = np.int32) + 3
            grid = StOptGrids.RegularLegendreGrid(lowValues, step, nbStep, poly)
             
            val = gasStorage(grid, maxLevelStorage)
            
            if world.rank == 0:
                self.assertAlmostEqual(val[0], val[1], None, None, accuracyClose)
              
if __name__ == '__main__':
    unittest.main()