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


# Generic test case
# p_nDim                   dimension of the problem
# p_nbSimul                number of simulations used
# p_level                  level of the sparse grid
# p_degree                 degree for the interpolation
# p_referenceValue         reference value
# p_accuracyEqual          accuracy for equality
def americanSparse(p_nDim, p_nbSimul, p_level, p_degree, p_referenceValue, p_accuracyEqual) :
    
    initialValues = np.zeros(p_nDim) + 1.
    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)
    # regressor
    weight = np.zeros(p_nDim) + 1.
    regressor = reg.SparseRegression(p_level, weight, p_degree)
    # bermudean value
    value = amer.resolution(simulator, payoff, regressor)
    return value
        
# test cases
class testAmericanOptionForSparse(unittest.TestCase):
    
    def test_americanSparseBasket1D(self):
        
        # dimension
        nDim = 1
        nbSimul = 100000
        level = 5
        referenceValue = 0.06031
        accuracyEqual = 0.2
    
        val = americanSparse(nDim, nbSimul, level, 1, referenceValue, accuracyEqual)
        self.assertAlmostEqual(val, referenceValue, None, None, accuracyEqual)
        val = americanSparse(nDim, nbSimul, level, 2, referenceValue, accuracyEqual)
        self.assertAlmostEqual(val, referenceValue, None, None, accuracyEqual)
        val = americanSparse(nDim, nbSimul, level, 3, referenceValue, accuracyEqual)
        self.assertAlmostEqual(val, referenceValue, None, None, accuracyEqual)
        
    def test_americanSparseBasket2D(self):
        
        # dimension
        nDim = 2
        nbSimul = 100000
        level = 5
        referenceValue = 0.03882
        accuracyEqual = 0.5
    
        val = americanSparse(nDim, nbSimul, level, 1, referenceValue, accuracyEqual)
        self.assertAlmostEqual(val, referenceValue, None, None, accuracyEqual)
        val = americanSparse(nDim, nbSimul, level, 2, referenceValue, accuracyEqual)
        self.assertAlmostEqual(val, referenceValue, None, None, accuracyEqual)
        val = americanSparse(nDim, nbSimul, level, 3, referenceValue, accuracyEqual)
        self.assertAlmostEqual(val, referenceValue, None, None, accuracyEqual)
        
    def test_americanSparseBasket3D(self):
        
        # dimension
        nDim = 3
        nbSimul = 100000
        level = 5
        referenceValue = 0.02947
        accuracyEqual = 0.5
    
        val = americanSparse(nDim, nbSimul, level, 1, referenceValue, accuracyEqual)
        self.assertAlmostEqual(val, referenceValue, None, None, accuracyEqual)
        val = americanSparse(nDim, nbSimul, level, 2, referenceValue, accuracyEqual)
        self.assertAlmostEqual(val, referenceValue, None, None, accuracyEqual)
        val = americanSparse(nDim, nbSimul, level, 3, referenceValue, accuracyEqual)
        self.assertAlmostEqual(val, referenceValue, None, None, accuracyEqual)
        
    def test_americanSparseBasket4D(self):
        
        # dimension
        nDim = 4
        nbSimul = 100000
        level = 5
        referenceValue = 0.02404
        accuracyEqual = 1.
    
        val = americanSparse(nDim, nbSimul, level, 1, referenceValue, accuracyEqual)
        self.assertAlmostEqual(val, referenceValue, None, None, accuracyEqual)
        val = americanSparse(nDim, nbSimul, level, 2, referenceValue, accuracyEqual)
        self.assertAlmostEqual(val, referenceValue, None, None, accuracyEqual)
        val = americanSparse(nDim, nbSimul, level, 3, referenceValue, accuracyEqual)
        self.assertAlmostEqual(val, referenceValue, None, None, accuracyEqual)
        
if __name__ == '__main__':
    
    unittest.main()
