# Copyright (C) 2016 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 utils.BasketOptions as bo
import simulators.BlackScholesSimulator as bs
import utils.americanOption as amer
import StOptGrids
import StOptReg as reg
import unittest



def american(p_nDim, p_nbSimul, p_nMesh, p_bRotate) :
    
    initialValues = np.zeros(p_nDim) + 1.0
    sigma = np.zeros(p_nDim) + 0.2
    mu = np.zeros(p_nDim) + 0.05
    corr = np.zeros((p_nDim,p_nDim))
    T = 1.
    nDate = 10
    np.fill_diagonal(corr, 1.)
    strike = 1.

    # simulator
    simulator = bs.BlackScholesSimulator(initialValues, sigma, mu, corr, T, nDate, p_nbSimul, False)
    # payoff
    payoff = bo.BasketPut(strike)
    # mesh
    nbMesh = np.zeros(p_nDim, dtype=np.int32) + p_nMesh
    # regressor
    regressor = reg.LocalLinearRegression(nbMesh,p_bRotate)
    # bermudean value
    value = amer.resolution(simulator, payoff, regressor)
    return value

# test cases
class testAmericanOptionTest(unittest.TestCase):
    
    def test_americanBasket1D(self):
  
        # dimension
        nDim = 1
        nbSimul = 500000
        nbMesh = 16
        referenceValue = 0.06031
        accuracyEqual = 0.13
               
        self.assertAlmostEqual(american(nDim, nbSimul, nbMesh, False), referenceValue, None, None, accuracyEqual)
        
    def test_americanBasket2D(self):    
        
        # dimension
        nDim = 2
        nbSimul = 1000000
        nbMesh = 16
        referenceValue = 0.03882
        accuracyEqual = 0.4
            
        self.assertAlmostEqual(american(nDim, nbSimul, nbMesh, False), referenceValue, None, None, accuracyEqual)
        self.assertAlmostEqual(american(nDim, nbSimul, nbMesh, True), referenceValue, None, None, accuracyEqual)
     
    def test_americanBasket3D(self):
    
        # dimension
        nDim = 3
        nbSimul = 2000000
        nbMesh = 10
        referenceValue = 0.02947
        accuracyEqual = 0.5
           
        self.assertAlmostEqual(american(nDim, nbSimul, nbMesh,False), referenceValue, None, None, accuracyEqual)
        self.assertAlmostEqual(american(nDim, nbSimul, nbMesh,True), referenceValue, None, None, accuracyEqual)
      
    def test_americanBasket4D(self):
        
        # dimension
        nDim = 4
        nbSimul = 2000000
        nbMesh = 6
        referenceValue = 0.02404
        accuracyEqual = 1.
           
        self.assertAlmostEqual(american(nDim, nbSimul, nbMesh,False), referenceValue, None, None, accuracyEqual)

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