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

accuracyClose = 0.2
accuracyEqual = 1e-10


class testSwingOptimSimuMpiTest(unittest.TestCase) :
    
    def test_resolutionSwingOptimSimuMpi(self):
    
        moduleMpi4Py=importlib.util.find_spec('mpi4py')
        if (moduleMpi4Py is not None):
            from mpi4py import MPI
            world = MPI.COMM_WORLD
            initialValues = np.zeros(1) + 1.
            sigma = np.zeros(1) + 0.2
            mu = np.zeros(1) + 0.05
            corr = np.ones((1,1))
            # number of step
            nStep = 20
            # exercise date
            dates = np.linspace(0., 1., nStep + 1)
            N = 3 # number of exercises
            T = 1.
            strike = 1.
            nbSimul = 100000
            nMesh = 8
            # payoff for swing
            payOffBasket = Utils.BasketCall(strike);
            payoff = Utils.PayOffSwing(payOffBasket,N)
            nbMesh = np.zeros(1, dtype = np.int32) + nMesh
            # Grid
            lowValues = np.zeros(1, dtype = float)
            step = np.zeros(1, dtype = float) + 1.
            nbStep = np.zeros(1, dtype = np.int32) + N
            grid =  StOptGrids.RegularSpaceGrid(lowValues, step, nbStep)
            # initial values
            initialStock = np.zeros(1)
            initialRegime = 0
            fileToDump = "CondExpSwingOptSimHLMpi"
            optimizer = opt.OptimizerSwingBlackScholes(payOffBasket, N)
            bOneFile = True
            # simulator
            simulatorBack = sim.BlackScholesSimulator(initialValues, sigma, mu, corr, dates[len(dates) - 1], len(dates) - 1, nbSimul, False)
            # regressor
            regressor = StOptReg.LocalLinearRegression(nbMesh)
            # link the simulations to the optimizer
            optimizer.setSimulator(simulatorBack)
            # bermudean value
            valueOptimMpi = dyn.DynamicProgrammingByRegressionDist(grid, optimizer, regressor, payoff, initialStock, initialRegime, fileToDump, bOneFile)
            world.barrier()
            # simulation value
            simulatorForward = sim.BlackScholesSimulator(initialValues, sigma, mu, corr, dates[len(dates) - 1], len(dates) - 1, nbSimul, True)
            optimizer.setSimulator(simulatorForward)
            valSimuOneFile = srt.SimulateRegressionDist(grid, optimizer, payoff, initialStock, initialRegime, fileToDump, bOneFile)
            print("Optim", valueOptimMpi, "SIMU", valSimuOneFile)
            
            if world.rank == 0:
                self.assertAlmostEqual(valueOptimMpi, valSimuOneFile, None, None, accuracyClose)
            
            optimizer = opt.OptimizerSwingBlackScholes(payOffBasket, N)
            bOneFile = False
            # simulator
            simulatorBack = sim.BlackScholesSimulator(initialValues, sigma, mu, corr, dates[len(dates) - 1], len(dates) - 1, nbSimul, False)
            # regressor
            regressor = StOptReg.LocalLinearRegression(nbMesh)
            # link the simulations to the optimizer
            optimizer.setSimulator(simulatorBack)
            # bermudean value
            valueOptimMpi = dyn.DynamicProgrammingByRegressionDist(grid, optimizer, regressor, payoff, initialStock, initialRegime, fileToDump, bOneFile)
             
            if world.rank == 0:
                print("Value in optimization", valueOptimMpi)
             
            world.barrier()
            # simulation value
            simulatorForward = sim.BlackScholesSimulator(initialValues, sigma, mu, corr, dates[len(dates) - 1], len(dates) - 1, nbSimul, True)
            optimizer.setSimulator(simulatorForward)
            valSimuMultipleFiles = srt.SimulateRegressionDist(grid, optimizer, payoff, initialStock, initialRegime, fileToDump, bOneFile)
            print("Optim", valueOptimMpi, "SIMU", valSimuMultipleFiles)
            self.assertAlmostEqual(valSimuOneFile, valSimuMultipleFiles, None, None, accuracyEqual)
            
if __name__ == '__main__':
    unittest.main()
