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

# unit test for global shape
############################

class OptimizerConstruction(unittest.TestCase):

    def test(self):
        # test MPI
        moduleMpi4Py=importlib.util.find_spec('mpi4py')
        if (moduleMpi4Py is not None):
            from mpi4py import MPI
        initialValues = np.zeros(1,dtype=float) + 1.
        sigma = np.zeros(1) + 0.2
        mu = np.zeros(1) + 0.05
        corr = np.ones((1,1),dtype=float)
        # number of step
        nStep = 30
        # exercise dates
        dates = np.linspace(0., 1., nStep + 1)
        T= dates[len(dates) - 1]
        nbSimul = 10 # simulation number (optimization and simulation)
        # simulator
        ##########
        bsSim = sim.BlackScholesSimulator(initialValues, sigma, mu, corr, T, len(dates) - 1, nbSimul, False)
        strike = 1.
        # Pay off
        payOff= Utils.BasketCall(strike)
        # optimizer
        ##########
        N = 3   # number of exercise dates
        swiOpt = opt.OptimizerSwingBlackScholes(payOff,N)
        # link simulator to optimizer
        swiOpt.setSimulator(bsSim)
        # archive
        ########
        ar = StOptGeners.BinaryFileArchive("Archive","w")
        # regressor
        ##########
        nMesh = np.array([1])
        regressor = StOptReg.LocalLinearRegression(nMesh)
        # Grid
        ######
        # low value for the meshes
        lowValues =np.array([0.],dtype=float)
        # size of the meshes
        step = np.array([1.],dtype=float)
        # number of steps
        nbStep = np.array([N], dtype=np.int32)
        gridArrival =  StOptGrids.RegularSpaceGrid(lowValues,step,nbStep)
        gridStart   =  StOptGrids.RegularSpaceGrid(lowValues,step,nbStep-1)
        # pay off function for swing
        ############################
        payOffBasket = Utils.BasketCall(strike);
        payoff = Utils.PayOffSwing(payOffBasket,N)
        dir(payoff)
        # final step
        ############
        asset =bsSim.getParticles()
        fin = StOptGlobal.FinalStepDP(gridArrival,1)
        values = fin.set( payoff,asset)
        # transition time step
        #####################
        # on step backward and get  asset
        asset = bsSim.stepBackwardAndGetParticles()
        # update regressor
        regressor.updateSimulations(0,asset)
        transStep = StOptGlobal.TransitionStepRegressionDP(gridStart,gridArrival,swiOpt)
        valuesNextAndControl=transStep.oneStep(values,regressor)
        transStep.dumpContinuationValues(ar,"Continuation",1,valuesNextAndControl[0],valuesNextAndControl[1],regressor)
        # simulate time step
        ####################
        nbSimul= 10
        vecOfStates =[] # state of each simulation
        for i in np.arange(nbSimul):
            # one regime, all with same stock level (dimension 2), same realization of simulation (dimension 3)
            vecOfStates.append(StOptGlobal.StateWithStocks(1, np.array([0.]) , np.zeros(1)))
        arRead = StOptGeners.BinaryFileArchive("Archive","r")
        simStep = StOptGlobal.SimulateStepRegression(arRead,1,"Continuation",gridArrival,swiOpt)
        phi = np.zeros((1,nbSimul))
        VecOfStateNext = simStep.oneStep(vecOfStates, phi)

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