# Copyright (C) 2016 EDF
# All Rights Reserved
# This code is published under the GNU Lesser General Public License (GNU LGPL)
import StOptGrids
import StOptSDDP
import StOptGlobal
import Utils
import SDDPSimulators as sim
import SDDPOptimizers as opt
import numpy as NP
import unittest
import math
import imp
import sddp.backwardForwardSDDP as bfSDDP # import of the function written in python

# unitest equivalent of testDemandSDDP : here MPI version
# High level python interface : at level of the backwardForwardSDDP c++ file
############################################################################
def demandSDDPFunc(p_sigD, p_sampleOptim ,p_sampleCheckSimul):

        maturity = 40
        nstep = 40;
        
        # optimizer parameters
        kappaD = 0.2; # mean reverting coef of demand
        spot = 3 ; # spot price

         # define a a time grid
        timeGrid = StOptGrids.OneDimRegularSpaceGrid(0., maturity / nstep, nstep)
        
        # periodicity factor
        iPeriod = 52;
        # average demande values
        demValues = []

        for i in list(range(nstep + 1)) :
            demValues.append(2. + 0.4 * math.cos((math.pi * i * iPeriod) / nstep))
 
        # define average demand
        demGrid = Utils.FutureCurve(timeGrid, demValues)
 
        initialState = demGrid.get(0.)*NP.ones(1)

        finCut = StOptSDDP.SDDPFinalCut(NP.zeros((2,1)))

        # here cuts are not conditional to an uncertainty
        nbMesh = NP.array([],NP.int32)
        nbUncertainties = 1;
        
        # backward simulator
        backwardSimulator = sim.SimulatorGaussianSDDP(nbUncertainties,p_sampleOptim)
        # forward simulator
        forwardSimulator = sim.SimulatorGaussianSDDP(nbUncertainties)
         
        # Create the optimizer
        optimizer = opt.OptimizeDemandSDDP(p_sigD, kappaD,  demGrid, spot, backwardSimulator, forwardSimulator)

        # optimisation dates
        dates = NP.linspace( 0., maturity,nstep + 1);

        # names for archive
        nameRegressor = "RegressorDemand";
        nameCut = "CutDemand";
        nameVisitedStates = "VisitedStateDemand";

        #  precision parameter
        nIterMax = 40
        accuracyClose =  1.
        accuracy = accuracyClose / 100.
        nstepIterations = 4; # check for convergence between nstepIterations step

        values = StOptSDDP.backwardForwardSDDP(optimizer,  p_sampleCheckSimul, initialState, finCut, dates,  nbMesh, nameRegressor, nameCut, nameVisitedStates, nIterMax,
                                                    accuracy, nstepIterations);

        print("Values " , values)
        return values


# unitest equivalent of testDemandSDDP : here low interface python version
# Low level python interface : use backwardForwardSDDP.py
##########################################################################
def demandSDDPFuncLowLevel(p_sigD, p_sampleOptim ,p_sampleCheckSimul):

        maturity = 40
        nstep = 40;
        
        # optimizer parameters
        kappaD = 0.2; # mean reverting coef of demand
        spot = 3 ; # spot price

         # define a a time grid
        timeGrid = StOptGrids.OneDimRegularSpaceGrid(0., maturity / nstep, nstep)
        
 
        # periodicity factor
        iPeriod = 52;
        # average demande values
        demValues = []

        for i in list(range(nstep + 1)) :
            demValues.append(2. + 0.4 * math.cos((math.pi * i * iPeriod) / nstep))
 
        # define average demand
        demGrid =Utils.FutureCurve(timeGrid, demValues)
 
        initialState = demGrid.get(0.)*NP.ones(1)

        finCut = StOptSDDP.SDDPFinalCut(NP.zeros((2,1)))

        # here cuts are not conditional to an uncertainty
        nbMesh = NP.array([],NP.int32)
        nbUncertainties = 1;
        
        # backward simulator
        backwardSimulator = sim.SimulatorGaussianSDDP(nbUncertainties,p_sampleOptim)
        # forward simulator
        forwardSimulator = sim.SimulatorGaussianSDDP(nbUncertainties)
         
        # Create the optimizer
        optimizer = opt.OptimizeDemandSDDP(p_sigD, kappaD,  demGrid, spot, backwardSimulator, forwardSimulator)

        # optimisation dates
        dates = NP.linspace( 0., maturity,nstep + 1);

        # names for archive
        nameRegressor = "RegressorDemand";
        nameCut = "CutDemand";
        nameVisitedStates = "VisitedStateDemand";

        #  precision parameter
        nIterMax = 40
        accuracyClose =  1.
        accuracy = accuracyClose / 100.
        nstepIterations = 4; # check for convergence between nstepIterations step

        values = bfSDDP.backwardForwardSDDP(optimizer,  p_sampleCheckSimul, initialState, finCut, dates,  nbMesh, nameRegressor,
                                     nameCut, nameVisitedStates, nIterMax,
                                     accuracy, nstepIterations);

        return values
  

class testDemandSDDP(unittest.TestCase):
    def testDemandSDDP1D(self):
        try:
            imp.find_module('mpi4py')
            found = True
        except:
            print("Not parallel module found ")
            found = False
    
        if found :
            from mpi4py import MPI
            world = MPI.COMM_WORLD

            sigD = 0.6 ;
            sampleOptim = 500;
            sampleCheckSimul = 500;
            
            values = demandSDDPFunc(sigD, sampleOptim ,sampleCheckSimul)

            if (world.rank==0):
                    print("Values is ",values)
        
    def testDemandSDDP1DLowLevel(self):
        sigD = 0.6 ;
        sampleOptim = 500;
        sampleCheckSimul = 500;
        demandSDDPFuncLowLevel(sigD, sampleOptim ,sampleCheckSimul)


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