# Copyright (C) 2019 EDF
# All Rights Reserved
# This code is published under the GNU Lesser General Public License (GNU LGPL)
import numpy as np
import math
import StOptGrids
import StOptTree
import StOptGlobal
import StOptGeners
import Utils
import Simulators
import unittest

# test cases
class testAmericanOptionTest(unittest.TestCase):
    
    def test_americanTree(self):
  
        # define a a time grid
        timeGrid = StOptGrids.OneDimRegularSpaceGrid(0., 2., 1)
        # future values
        futureGrid =  Utils.FutureCurve(timeGrid, [50.,50.])

        # degenerate trinomial tri
        mr = 0.00001
        sigma =0.2

        # interest rate
        r = 0.05

        # number of time steps
        nstep =50
        
        # maturity
        T =1.
        
        # step
        dt = T / nstep
        
        # simulation dates
        dates =  dt * np.arange(0,nstep+1)

        # simulation dates
        treeSim = Simulators.TrinomialTreeOUSimulator(mr, sigma, dates)

        indexT = np.arange(0,nstep+1, 5)
        treeSim.dump("Tree",indexT)

        # read archive
        archiveToRead = StOptGeners.BinaryFileArchive("Tree", "r");
        
       # FIRST VALORIZATION
        ####################
        
        # backward simulator
        backSimulator = Simulators.MeanRevertingSimulatorTree(archiveToRead, futureGrid, sigma, mr)

        # strike
        strike  = 50.

        # actu
        actu = np.exp(r*dates[indexT[-1]])
        # spot : add interest rate
        spot = backSimulator.getSpotValues()*actu
        # actualized payoff
        val1= np.where( strike-spot>0,strike-spot,0)/actu
        for istep in np.arange(np.shape(indexT)[0]-1):
            # backward in simulator
            backSimulator.stepBackward()
            # get back probability matrix
            proba = backSimulator.getProba()
            # and connection matrix
            connected = backSimulator.getConnected()
            # creta tree for conditional expectation
            tree=StOptTree.Tree(proba,connected)      
            # interest rates
            actu = np.exp(r*dates[indexT[-2-istep]])
            # spot : add interest rate
            spot = backSimulator.getSpotValues()*actu
            # pay off 
            payOff = np.where( strike-spot>0,strike-spot,0)/actu
            # actualize value
            val1 = tree.expCond(val1)
            # arbitrage
            val1 = np.where( val1> payOff, val1, payOff)

        final = val1[0]
        
        indexT = np.arange(0,nstep+1)
        treeSim.dump("Tree",indexT)

        # read archive
        archiveToRead = StOptGeners.BinaryFileArchive("Tree", "r");
        
        # backward simulator
        backSimulator = Simulators.MeanRevertingSimulatorTree(archiveToRead, futureGrid, sigma, mr)

        # SECOND VALORIZATION
        ####################
        # actu
        actu = np.exp(r*dates[-1])
        # spot : add interest rate
        spot = backSimulator.getSpotValues()*actu
        # actualized payoff
        val2= np.where( strike-spot>0,strike-spot,0)/actu

        for istep in np.arange(nstep):
            # backward in simulator
            backSimulator.stepBackward()
            # get back probability matrix
            proba = backSimulator.getProba()
            # and connection matrix
            connected = backSimulator.getConnected()
            # creta tree for conditional expectation
            tree=StOptTree.Tree(proba,connected)      
            # actualize value
            val2 = tree.expCond(val2)

            if ((istep+1)%5==0):
                # interest rates
                actu = np.exp(r*dates[-2-istep])
                # spot : add interest rate
                spot = backSimulator.getSpotValues()*actu
                # pay off 
                payOff = np.where( strike-spot>0,strike-spot,0)/actu
                # arbitrage
                val2 = np.where( val2> payOff, val2, payOff)

        print("val1 " , val1 , " val2 " , val2)
        self.assertAlmostEqual(val1, val2, None, None, 0.000001)
        
 
if __name__ == '__main__':
    unittest.main()
