# Copyright (C) 2016, 2018 EDF
# All Rights Reserved
# This code is published under the GNU Lesser General Public License (GNU LGPL)
import numpy as np
import math as maths
import random
import unittest
import math


# Implement a Black Scholes simulator
class BlackScholesSimulator :
    
    # Constructor for Black Scholes simulator
    # p_initialValues  initial values for assets
    # p_sigma          volatility of assets
    # p_mu             trend for assets
    # p_correl         correlation for assets
    # p_T              maturity
    # p_nbStep         number of time step
    # p_nbSimul        number of simulations
    # p_bForward       true if the simulation is forward, false if backward
    def __init__(self, p_initialValues, p_sigma, p_mu, p_correl, p_T, p_nbStep, p_nbSimul, p_bForward) :
        
        self.m_initialValues = p_initialValues
        self.m_sigma = p_sigma
        self.m_mu = p_mu
        self.corrFactTrans = np.linalg.cholesky(p_correl)
        self.m_T = p_T
        self.m_step = p_T / p_nbStep
        self.m_nbStep = p_nbStep
        self.m_nbSimul = p_nbSimul
        self.m_bForward = p_bForward
        self.m_currentStep = 0. if p_bForward else p_T # Current step during resolution
        self.m_brownian = np.zeros((len(p_initialValues), p_nbSimul)) # store the brownian motion values (no correlation)
        
        np.random.seed(0)
       
        if self.m_bForward :
            self.m_brownian = np.zeros((len(p_initialValues), p_nbSimul))
        else :
            sqrtT = maths.sqrt(self.m_T)
            self.m_brownian = sqrtT * np.random.randn(len(p_initialValues), p_nbSimul)
            
    # gets brownian motions correlated and send back asset values
    def brownianToAsset(self) :
        
        assetToReturn = np.zeros((len(self.m_initialValues), self.m_nbSimul))
        firstTerm = (self.m_mu - 0.5 * self.m_sigma *  self.m_sigma) * self.m_currentStep
        assetToReturn[:,:] = (self.m_initialValues * np.exp(firstTerm + self.m_sigma * self.m_brownian.transpose())).transpose()
         
        return assetToReturn
        
    # a step forward for brownians
    def forwardStepForBrownian(self) :
        
        normalSample = np.random.randn(len(self.m_initialValues), self.m_nbSimul)
        self.m_brownian = (np.multiply(maths.sqrt(self.m_currentStep), normalSample.transpose())).transpose()  # Check shape normalSample =  shape noise

    # a step backward for brownians
    def backwardStepForBrownian(self) :
        
        if abs(self.m_currentStep - 0.) <= (0.01 * abs(0. + self.m_currentStep) * 10) :
            self.m_brownian = np.zeros((len(self.m_initialValues), self.m_nbSimul))
        else :
            util1 = max(self.m_currentStep / (self.m_currentStep + self.m_step), 0.)
            util2 = maths.sqrt(util1 * self.m_step)
            # use brownian bridge
            self.m_brownian = (np.multiply(util1, self.m_brownian) + np.multiply(util2, np.random.randn(len(self.m_initialValues), self.m_nbSimul)))
            
    # get current asset values
    def getParticles(self) :
        
        return self.brownianToAsset()

    # a step forward for simulations
    def stepForward(self) :
        
        if self.m_bForward == False :
            pass
        
        self.m_currentStep += self.m_step
        self.forwardStepForBrownian()

    # return the asset values (asset,simulations)
    def stepBackward(self) :
        
        if self.m_bForward == True :
            pass
        
        self.m_currentStep -= self.m_step
        self.backwardStepForBrownian()

    # a step forward for simulations
    # return the asset values (asset,simulations)
    def stepForwardAndGetParticles(self) :
        
        if self.m_bForward == False :
            pass
        
        self.m_currentStep += self.m_step
        self.forwardStepForBrownian()
        
        return self.brownianToAsset()

    # a step backward for simulations
    # return the asset values (asset,simulations)
    def stepBackwardAndGetParticles(self) :
        
        if self.m_bForward == True :
            pass
        
        self.m_currentStep -= self.m_step
        self.backwardStepForBrownian()
        
        return self.brownianToAsset()

    # Get back attribute
    
    def getCurrentStep(self) :
        
        return self.m_currentStep

    def getT(self) :
       
        return self.m_T

    def getStep(self) :
        
        return self.m_step

    def getInitialValues(self) :
        
        return self.m_initialValues

    def getSigma(self) :
        
        return self.m_sigma

    def getMu(self) :
        
        return self.m_mu

    def getNbSimul(self) :
        
        return self.m_nbSimul

    def getNbStep(self) :
        
        return self.m_nbStep

    def getDimension(self) :
        
        return len(self.m_sigma)

    # actualize at date t=0
    def getActu(self):
        return math.exp(- self.m_mu[0]* self.m_currentStep )

    # actualize on one step
    def getActuStep(self):
        return math.exp(- self.m_mu[0]* self.m_step )
    
class BlackScholesSimulatorTest(unittest.TestCase):
    
    def test_callOption(self):
        
        S_0 = np.zeros(2) + 100.
        vol = np.zeros(2) + 0.25
        mu = np.zeros(2) + 0.1
        corr = np.eye(2)
        T = 2.
        nbStep = 10
        nbSimul = 100000
        K = 100.
        
        ga = BlackScholesSimulator(S_0, vol, mu, corr, T, nbStep, nbSimul, True)
        bu = BlackScholesSimulator(S_0, vol, mu, corr, T, nbStep, nbSimul, False)
        SF = np.zeros(nbStep / 2)
        SB = np.zeros(nbStep / 2)
        
        def compareList(l, res) :
            for i in list(range(len(l))) :
                res[i] = max(l[i], 0.)
            
            return res
     
        for i in range(nbStep / 2) :
            lF = np.zeros(nbSimul)
            compareList(ga.stepForwardAndGetParticles()[0,:] - K,lF)
            SF[i] = np.mean(lF * maths.exp(-mu[0] * T/2))
            lB = np.zeros(nbSimul)
            compareList(bu.stepBackwardAndGetParticles()[0,:] - K,lB)
            SB[i] = np.mean(lB * maths.exp(-mu[0] * T/2))
             
        self.assertAlmostEqual(SF[nbStep / 2 - 1], 15, None, None, 0.2)
        self.assertAlmostEqual(SB[nbStep / 2 - 1], 15, None, None, 0.2)
        self.assertAlmostEqual(SF[nbStep / 2 - 1], SB[nbStep / 2 - 1], None, None, 0.2)
          
    def test_martingale(self):
        
        S_0 = np.zeros(2) + 100.0
        vol = np.zeros(2) + 0.25
        mu = np.zeros(2)
        corr = np.eye(2)
        T = 10.
        nbStep = 10
        nbSimul = 1000000
        K = 100.
        dp = np.zeros(nbStep)
          
        zor = BlackScholesSimulator(S_0, vol, mu, corr, T, nbStep, nbSimul, True)
          
        for j in range(nbStep) :
            dp[j] = np.mean(zor.stepForwardAndGetParticles()[0,:] * maths.exp(-mu[0] * (j + 1)))
            self.assertAlmostEqual(dp[j], S_0[0], None, None, 0.2)
            
if __name__ == '__main__':
    unittest.main()
