#!/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 utils.OneDimRegularSpaceGrid as rsg
import utils.OneDimData as data
import simulators.MeanRevertingSimulator as mrsim
import dp.OptimizeGasStorage as ogs
import StOptReg as reg
import StOptGrids
import dp.DynamicProgrammingByRegressionVaryingGrids as dyn
import unittest

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_timeChangeGrid, p_grids, p_maxLevelStorage) :
    
    # storage
    injectionRateStorage = 60000
    withdrawalRateStorage = 45000
    injectionCostStorage = 0.35
    withdrawalCostStorage = 0.35

    maturity = 1.
    nstep = 100

    # define a a time grid
    timeGrid = rsg.OneDimRegularSpaceGrid(np.zeros(1), maturity / nstep, nstep)
    # future values
    futValues = np.zeros(nstep + 1)
    
    # periodicity factor
    iPeriod = 52
    
    for i in list(range(nstep + 1)) :
        futValues[i] = 50. + 20 * math.sin((math.pi * i * iPeriod) / nstep)
        
    # define the future curve
    futureGrid = data.OneDimData(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 = mrsim.MeanRevertingSimulator(futureGrid, sigma, mr, rate, maturity, nstep, nbsimulOpt, bForward)
    # optimizer
    storage = ogs.OptimizeGasStorage(injectionRateStorage, withdrawalRateStorage, injectionCostStorage, withdrawalCostStorage)
    
    # regressor
    nMesh = 6
    nbMesh = np.zeros(1, dtype = np.int32) + nMesh
    regressor = reg.LocalLinearRegression(nbMesh)
    # final value
    vFunction = ZeroFunction()
    
    # initial values
    initialStock = np.zeros(1) + p_maxLevelStorage
    initialRegime = 0 # only one regime
    
    # Optimize
    fileToDump = "CondExpGasStorageVaryingCavity"
    
    # link the simulations to the optimizer
    storage.setSimulator(backSimulator)
    valueOptim = dyn.DynamicProgrammingByRegression(p_timeChangeGrid, p_grids, storage, regressor, vFunction, initialStock, initialRegime, fileToDump)
    print("valueOptim", valueOptim)
            
class testGasStorageVaryingCavityTest(unittest.TestCase):
          
    def test_simpleStorageVaryingCavity(self):
              
        # storage
        maxLevelStorage = 90000
        # grid
        timeChangeGrid = []
        grids = []
        nGrid = 10
        
        lowValues1 = np.zeros(1)
        step1 = np.zeros(1) + (maxLevelStorage / nGrid)
        nbStep1 = np.zeros(1, dtype = np.int32) + nGrid
        timeChangeGrid.append(0.)
        grids.append(StOptGrids.RegularSpaceGrid(lowValues1, step1, nbStep1))
        
        lowValues2 = np.zeros(1) + 30000.
        step2 = np.zeros(1) + 10000.
        nbStep2 = np.zeros(1, dtype = np.int32) + 3
        timeChangeGrid.append(0.3)
        grids.append(StOptGrids.RegularSpaceGrid(lowValues2, step2, nbStep2))
        
        lowValues3 = np.zeros(1)
        step3 = np.zeros(1) + 15000.
        nbStep3 = np.zeros(1, dtype = np.int32) + 6
        timeChangeGrid.append(0.7)
        grids.append(StOptGrids.RegularSpaceGrid(lowValues3, step3, nbStep3))
        
        gasStorage(timeChangeGrid, grids, maxLevelStorage)
              
if __name__ == '__main__':
    
    unittest.main()