# Copyright (C) 2021 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.DynamicProgrammingSwitchingByRegressionHighLevel  as dyn
import dp.SimulateSwitchingRegression  as simOptCont
import unittest
import importlib


# valorisation of a thermal asset
    
class testThermalASsetHighLevelTest(unittest.TestCase):

    def testThermalAsset(self):
        # test MPI
        moduleMpi4Py=importlib.util.find_spec('mpi4py')
        if (moduleMpi4Py is not None):
            from mpi4py import MPI
        T = 100.
        # number of time step  ( step of 4 hours : 6 step per day)
        nStepPerday = 6
        nbDays = int(T)
        nstep = nStepPerday * nbDays;
        # define a a time grid
        timeGrid = StOptGrids.OneDimRegularSpaceGrid(0., T / nstep, nstep)
  
        #to store the two assets
        futGrids = []
        futValues =  [ 30. + 5 * np.cos((2 * math.pi * i) / nstep) + np.cos(2 * math.pi * i / (7.*nStepPerday)) for i in range(nstep+1)]
        futGrids.append(Utils.FutureCurve(timeGrid, futValues))
        futValues = [ 30. + 5 * np.cos((2 * math.pi * i) / nstep) + np.cos(2 * math.pi* i / (7.*nStepPerday))  + 5.*np.cos(2 * math.pi * i / nStepPerday)   for i in range(nstep+1)]    
        futGrids.append(Utils.FutureCurve(timeGrid, futValues))
        
        sigma = [0.02,0.08] # volatility by product
        mr = [0.004,0.01] # mean reverting by product
        correl = np.identity(2)
        correl[0, 1] = 0.8
        correl[1, 0] = 0.8

        # Constraints
        nbMinOn = 6  #  minimal number of time step the  thermal asset is on
        nbMinOff = 4 #  maximal number of time step the terhma asset is off

        # grids
        gridPerReg =  [ StOptGrids.RegularSpaceIntGrid(np.zeros([1]), (nbMinOn-1)*np.ones([1])),
                        StOptGrids.RegularSpaceIntGrid(np.zeros([1]), (nbMinOff-1)*np.ones([1]))]         

        
        iRegInit = 0 # on at begininng
        pointState = np.zeros([1]) # on since previous time step
        
         # create Thermal Asset object
        switchOffToOnCost = 4
        
        
        # optimizer
        thermalAsset = opt.OptimizeThermalAssetMeanReverting1DAssets(switchOffToOnCost)

        # simulator
        nbSimulOpt =5000
        iSeed = 1
        simulator = sim.MeanReverting1DAssetsSimulator(futGrids, mr, sigma, correl, T , nstep, nbSimulOpt, False,"ArchiveSim", iSeed)

        # regressor
        nMesh = 2
        nbMesh = np.zeros(2, dtype = np.int32) + nMesh
        regressor = reg.LocalLinearRegression(nbMesh)

        # affect simulator to optimizer
        thermalAsset.setSimulator(simulator)
        
        # optimize
        valOptim = dyn.DynamicProgrammingSwitchingByRegressionHighLevel(gridPerReg, thermalAsset, regressor, pointState,iRegInit, "BellmanThermal")
        
        print("valOptimp", valOptim)

        # revert simulator
        iSeed = 10
        nbSimulForward = 5000
        simulator = sim.MeanReverting1DAssetsSimulator(futGrids, mr, sigma, correl, T , nstep, nbSimulForward, True,"ArchiveSim", iSeed)
      
        # affect simulator to optimizer
        thermalAsset.setSimulator(simulator)

        # simulate
        valSimu = simOptCont.SimulateSwitchingRegression(gridPerReg,thermalAsset,pointState,iRegInit, "BellmanThermal")
        print("val simulation", valSimu)

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