# 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 dynmpi
import dp.SimulateRegressionDist as srtmpi
import importlib

accuracyClose = 1.5


# p_ndim        dimension of the swing
# p_bOneFile    Do we use one unique file for continuation values
def swingND(p_ndim, p_bOneFile):
    
    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 = 80000
    nMesh = 8
    # payoff for swing
    payOffBasket = Utils.BasketCall(strike)
    payoff = Utils.PayOffFictitiousSwing(payOffBasket,N)
    nbMesh = np.zeros(1, dtype = np.int32) + nMesh
    # simulator
    simulator = sim.BlackScholesSimulator(initialValues, sigma, mu, corr, dates[len(dates) - 1], len(dates) - 1, nbSimul, False)
    # Grid
    lowValues = np.zeros(p_ndim)
    step = np.zeros(p_ndim) + 1.
    # the stock is discretized with values from 0 to N included
    nbStep = np.zeros(p_ndim, dtype = np.int32) + N
    grid =  StOptGrids.RegularSpaceGrid(lowValues, step, nbStep)
    # optimizer
    optimizer = opt.OptimizerFictitiousSwingBlackScholes(payOffBasket, N, p_ndim)
    # initial values
    initialStock = np.zeros(p_ndim)
    initialRegime = 0
    fileToDump = "CondExpSwingOptimNDHLMpi"
    # regressor
    regressor = StOptReg.LocalLinearRegression(nbMesh)
    # link the simulations to the optimizer
    optimizer.setSimulator(simulator)
    # bermudean value
    valueOptimMpi = dynmpi.DynamicProgrammingByRegressionDist(grid, optimizer, regressor, payoff, initialStock, initialRegime, fileToDump, p_bOneFile)
    # simulation value
    simulatorForward = sim.BlackScholesSimulator(initialValues, sigma, mu, corr, dates[len(dates) - 1], len(dates) - 1, nbSimul, True)
    optimizer.setSimulator(simulatorForward)
    valSimuMpi = srtmpi.SimulateRegressionDist(grid, optimizer, payoff, initialStock, initialRegime, fileToDump, p_bOneFile)
    print("Optim", valueOptimMpi, "SIMU", valSimuMpi)
    return valueOptimMpi, valSimuMpi

class testSwingOptimSimuNDMpiTest(unittest.TestCase):
    
    def test_swingOptionOptim2DSimuMpiOneFile(self):
        
       moduleMpi4Py=importlib.util.find_spec('mpi4py')
       if (moduleMpi4Py is not None):
           from mpi4py import MPI
           world = MPI.COMM_WORLD
           ndim = 2
           bOneFile = True
           val = swingND(ndim, bOneFile)
           
           if world.rank == 0:
               self.assertAlmostEqual(val[0], val[1], None, None, accuracyClose)
        
    def test_swingOptionOptim2DSimuMpiMultipleFile(self):
        
      moduleMpi4Py=importlib.util.find_spec('mpi4py')
      if (moduleMpi4Py is not None):
          from mpi4py import MPI
          world = MPI.COMM_WORLD
          ndim = 2
          bOneFile = False
          val = swingND(ndim, bOneFile)
          
          if world.rank == 0:
              self.assertAlmostEqual(val[0], val[1], None, None, accuracyClose)
            
if __name__ == '__main__':
    
    unittest.main()
