# 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 unittest
import random
import StOptGrids
import StOptReg
import StOptTree
import StOptGlobal
import Simulators
import StOptGeners


# unit test for tree

class testTree(unittest.TestCase):

    def testContinuation(self):

        # Mean Reverting model
        mr = 0.3;
        sig = 0.6;

        # nb grid points
        nbStock=4
        
        # step
        dt = 1. / 100.
        
        # simulation dates
        dates =  dt * np.arange(0,16)

        # simulaton dates
        tree = Simulators.TrinomialTreeOUSimulator(mr, sig, dates)

        iFirst = 10
        iLast = 14

        # nodes at dates 5
        points = tree.getPoints(iFirst)

        # nodes at last date
        pointsNext = tree.getPoints(iLast)

        # probabilities
        proba = tree.getProbability(iFirst, iLast)

        # connection matrix
        connectAndProba =  tree.calConnected(proba);

        # to regress
        toTreeress= np.zeros( (nbStock, np.shape(pointsNext)[1]))
        for i in range(nbStock):
            toTreeress[i,:] = i + pointsNext[0,:]

        # grid for storage
        grid = StOptGrids.RegularSpaceGrid(np.array([0.]), np.array([1.]), np.array([nbStock - 1]))

        # conditional expectation object by trees
        tree= StOptTree.Tree(connectAndProba[1],connectAndProba[0])

        # conditional expectation taken
        valTree = tree.expCondMultiple(toTreeress)

        # continuation object
        continuation = StOptTree.ContinuationValueTree(grid, tree,  toTreeress.transpose())

        # interpolation point
        ptStock = np.array([0.5*nbStock])

        #  conditional expectation using continuation object
        treeByContinuation = continuation.getValueAtNodes(ptStock);

        for i in range(np.shape(points)[1]):
            self.assertAlmostEqual(treeByContinuation[i],valTree[int(nbStock / 2), i], 7,"Difference between function and its condition expectation estimated greater than tolerance")
 
        # into a list  : two times to test 2 regimes
        listToReg = []
        listToReg.append(toTreeress.transpose())
        listToReg.append(toTreeress.transpose())
 
        #default non compression
        archiveToWrite = StOptGeners.BinaryFileArchive("MyArchive", "w");
        # continuation value dump
        archiveToWrite.dumpGridTreeValue("toStore",1,listToReg,tree,grid)

        # read archive
        archiveToRead =  StOptGeners.BinaryFileArchive("MyArchive","r")
        contValues = archiveToRead.readGridTreeValue(1,"toStore")
        # interpolate
        value =contValues[0].getValues(ptStock)
        for i in range(np.shape(points)[1]):
            self.assertAlmostEqual(value[i],valTree[int(nbStock / 2), i], 7,"Difference between function and its condition expectation estimated greater than tolerance")
  
if __name__ == '__main__': 
    unittest.main()
