# 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 simulators.BlackScholesSimulator as bs
import utils.BasketOptions as bo
import StOptGrids 
import StOptReg as reg
import StOptGlobal
import Simulators as sim
import Optimizers as opt
import Utils
from utils.EuropeanOptions import CallOption
from utils.EuropeanOptions import PutOption as put
import importlib
import dp.DynamicProgrammingByRegressionDist as dynmpi

accuracyClose = 0.2
accuracyNearlyEqual = 0.05


# Analytical value
# p_S       asset value
# p_sigma   volatility
# p_r       interest rate
# p_strike  strike
# p_dates   possible exercise dates
# return option value
def analyticalValue(N, p_S, p_sigma, p_r, p_strike, p_dates) :
    
    analytical = 0.
    callOption = CallOption()
    
    for i in range(len(p_dates) - N, len(p_dates)) :
        analytical += callOption.operator(p_S, p_sigma, p_r, p_strike, p_dates[i])

    return analytical

# Classical resolution for swing
# p_sim             Monte Carlo simulator
# p_payOff          Option pay off
# p_regressor       regressor object
# p_dates           possible exercise dates
# p_N               number of exercises
def resolutionSwing(p_sim, p_payOff, p_regressor, p_dates, p_N) :
    
    if ((p_sim.getNbStep() + 1) == len(p_dates)) == False :
        pass
    
    # asset simulated under the neutral risk probability : get the trend of first asset to get interest rate
    # in this example the step between two exercises is given
    expRate = p_sim.getActuStep()
    
    # final payOff
    finalPayOff = p_payOff.applyVec(p_sim.getParticles())
    # Terminal function depending on simulations and stock already exercised
    cashNext = np.zeros((len(finalPayOff), p_N))

    for i in range(len(finalPayOff)) :
        cashNext[i,:] = (finalPayOff[i]).transpose()

    cashPrev = np.zeros((len(finalPayOff), p_N))

    for iStep in range(len(p_dates) - 1)[::-1] :
        asset = p_sim.stepBackwardAndGetParticles()
        payOffLoc = p_payOff.applyVec(asset)
        payOffLoc = payOffLoc.squeeze()

        # conditional expectation
        if iStep == 0 :
            p_regressor.updateSimulations(True, asset)
        else :
            p_regressor.updateSimulations(False, asset)

        # store conditional expectations
        vecCondEspec = list(range(p_N))

        for iStock in range(p_N):
            vecCondEspec[iStock] = p_regressor.getAllSimulations(cashNext[:,iStock]).squeeze() * expRate

        # arbitrage
        for iStock in range(p_N - 1):
            cashPrev[:,iStock] = np.where(payOffLoc + vecCondEspec[iStock + 1] > vecCondEspec[iStock], payOffLoc + expRate * cashNext[:,iStock + 1], expRate * cashNext[:,iStock])
            
        # last stock
        cashPrev[:,p_N - 1] = np.where(payOffLoc > vecCondEspec[p_N - 1], payOffLoc, expRate * cashNext[:,p_N - 1])
            
        # switch
        tempVec = cashNext
        cashNext = cashPrev
        cashPrev = tempVec

    return np.mean(cashNext[:,0])

def resolutionSwingContinuation(p_sim, p_payOff, p_regressor, p_dates, p_N):
    
    if (p_sim.getNbStep() + 1 == len(p_dates)) == False:
        pass
    
    # asset simulated under the neutral risk probability : get the trend of first asset to get interest rate
    # in this example the step between two exercises is given
    expRate = math.exp(-p_sim.getStep() * p_sim.getMu()[0])
    lowValues = np.zeros(1)
    step = np.zeros(1)
    lowValues[0] = 0.
    step[0] = 1
    nbStep = np.zeros(1, dtype = np.int32)
    nbStep[0] = p_N - 1
    regular = StOptGrids.RegularSpaceGrid(lowValues, step, nbStep)
    extremeGrid = np.zeros(2)
    extremeGrid[0] = lowValues
    extremeGrid[1] = lowValues + step * nbStep
    
    # final payOff
    finalPayOff = p_payOff.applyVec(p_sim.getParticles())
    # Terminal function depending on simulations and stock already exercised
    cashNext = np.zeros((len(finalPayOff), regular.getNbPoints()))

    for i in range(len(finalPayOff)):
        cashNext[i,:] = (np.zeros(regular.getNbPoints()) + finalPayOff[i]).transpose()

    cashPrev = np.zeros((len(finalPayOff), regular.getNbPoints()))

    for iStep in range(len(p_dates) - 2)[::-1] :
        asset = p_sim.stepBackwardAndGetParticles()
        payOffLoc = p_payOff.applyVec(asset)

        # conditional expectation
        if iStep == 0 :
            p_regressor.updateSimulations(True, asset)
        else :
            p_regressor.updateSimulations(False, asset)
            
        # continuation value object dealing with stocks
        continuation = reg.ContinuationValue(regular, p_regressor, cashNext)
        # iterator on grid points
        iterOnGrid = regular.getGridIterator()
        
        while iterOnGrid.isValid():
            CoordStock = iterOnGrid.getCoordinate()
            # use continuation to get realization of condition expectation
            conditionExpecCur = np.multiply(expRate, continuation.getAllSimulations(CoordStock).squeeze())
            
            if CoordStock[0] + 1 <= extremeGrid[1] :
                conditionExpecNext = np.multiply(expRate, continuation.getAllSimulations(CoordStock + 1).squeeze())
                cashPrev[:,iterOnGrid.getCount()] = np.where(payOffLoc + conditionExpecNext > conditionExpecCur, payOffLoc + np.multiply(expRate, cashNext[:,iterOnGrid.getCount() + 1]), np.multiply(expRate, cashNext[:,iterOnGrid.getCount()]))
            else :
                cashPrev[:,iterOnGrid.getCount()] = np.where(payOffLoc > conditionExpecCur, payOffLoc, np.multiply(expRate, cashNext[:,iterOnGrid.getCount()]))
                
            iterOnGrid.next()
            
        # switch pointer
        tempVec = cashNext
        cashNext = cashPrev
        cashPrev = tempVec
        
    return np.mean(cashNext[:,0])

#  function to test stock in dimension above 1
def multiStock(p_ndim):
    
    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
        # exercice date
        dates = np.linspace(0., 1., nStep + 1)
        # exercice dates
        N = 3
        T = 1.
        strike = 1.
        nbSimul = 10000
        nMesh = 4
        # payoff
        payoff = Utils.BasketCall(strike)
        # store sequential
        valueSeq = 0
        # mesh
        nbMesh = np.zeros(1, dtype = np.int32) + nMesh
        
        if world.rank == 0:
            # simulator
            simulator = bs.BlackScholesSimulator(initialValues, sigma, mu, corr, dates[len(dates) - 1], len(dates) - 1, nbSimul, False)
            # regressor
            regressor = reg.LocalLinearRegression(nbMesh)
            # bermudean value
            valueSeq = resolutionSwing(simulator, payoff, regressor, dates, N)
            
        # 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)
        # final value
        vFunction = Utils.PayOffFictitiousSwing(payoff, N)
        # optimizer
        optimizer = opt.OptimizerFictitiousSwingBlackScholes(payoff, N, p_ndim)
        # initial values
        initialStock = np.zeros(p_ndim)
        initialRegime = 0
        fileToDump = "CondExpSwing"
        # regressor
        regressor = reg.LocalLinearRegression(nbMesh)
        bOneFile = False
        # link the simulations to the optimizer
        optimizer.setSimulator(simulator)
        valueParal = dynmpi.DynamicProgrammingByRegressionDist(grid, optimizer, regressor, vFunction, initialStock, initialRegime, fileToDump, bOneFile)
        return valueParal, valueSeq

class testSwingOptionTest(unittest.TestCase) :
     
    def test_swingOptionInOptimization(self):
         
        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 = 30
        # exercise date
        dates = np.linspace(0., 1., nStep + 1)
        N = 3 # 3 exercise dates
        T = 1.
        strike = 1.
        nbSimul = 200000
        nMesh = 16
        # payoff
        payoff = bo.BasketCall(strike)
        # analytical
        analytical = analyticalValue(N, initialValues[0], sigma[0], mu[0], strike, dates)
        
        # mesh
        nbMesh = np.zeros(1, dtype = np.int32) + nMesh
        # simulator
        simulator = bs.BlackScholesSimulator(initialValues, sigma, mu, corr, dates[len(dates) - 1], len(dates) - 1, nbSimul, False)
        # regressor
        regressor = reg.LocalLinearRegression(nbMesh)
        # bermudean value
        valueSeq = resolutionSwing(simulator, payoff, regressor, dates, N)
        
        # simulator
        simulator = bs.BlackScholesSimulator(initialValues, sigma, mu, corr, dates[len(dates) - 1], len(dates) - 1, nbSimul, False)
        # regressor
        regressor = reg.LocalLinearRegression(nbMesh)
        # using continuation values
        valueSeqContinuation = resolutionSwingContinuation(simulator, payoff, regressor, dates, N)
        print("valueSeq", valueSeq, "valueSeqContinuation", valueSeqContinuation)
        
        self.assertAlmostEqual(valueSeq, analytical, None, None, accuracyClose)
        
    def test_swingOption2D(self):
        
        moduleMpi4Py=importlib.util.find_spec('mpi4py')
        if (moduleMpi4Py is not None):
            from mpi4py import MPI
            world = MPI.COMM_WORLD
            val = multiStock(2)
            
            if world.rank == 0:
                self.assertAlmostEqual(val[1] * 2, val[0], None, None, accuracyNearlyEqual)
                
    def test_swingOption3D(self):
        
        moduleMpi4Py=importlib.util.find_spec('mpi4py')
        if (moduleMpi4Py is not None):
            from mpi4py import MPI
            world = MPI.COMM_WORLD
            val = multiStock(3)
            
            if world.rank == 0:
                self.assertAlmostEqual(val[1] * 3, val[0], None, None, accuracyNearlyEqual)
                
if __name__ == '__main__':
    
    unittest.main()
