# 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.DynamicProgrammingByRegressionHighLevel as dyn
import dp.SimulateRegressionControlHighLevel as srt
import imp

accuracyClose = 0.25


class testSwingOptionTest(unittest.TestCase) :

    def test_resolutionSwing(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
            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 = 4
            # exercise date
            dates = np.linspace(0., 1., nStep + 1)
            N = 2 # 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)
            step = np.zeros(1) + 1.
            nbStep = np.zeros(1, dtype = np.int32) + N
            grid =  StOptGrids.RegularSpaceGrid(lowValues, step, nbStep)
            # optimizer
            optimizer = opt.OptimizerSwingBlackScholes(payOffBasket, N)
            # initial values
            initialStock = np.zeros(1)
            initialRegime = 0
            fileToDump = "CondExpSwingOptSimHL"
            # 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
            valueOptim = dyn.DynamicProgrammingByRegressionHighLevel(grid, optimizer, regressor, payoff, initialStock, initialRegime, fileToDump)
            # simulation value
            print("Optim", valueOptim)

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