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


# unit test for dumping binary archive of regressed value and Read then
######################################################################

class testBinaryArchiveStOpt(unittest.TestCase):

      
      def testSimpleStorageAndLectureRecGrid(self):
        
        # low value for the mesh
        lowValues =np.array([1.,2.,3.],dtype=np.float64)
        # size of the mesh
        step = np.array([0.7,2.3,1.9],dtype=np.float64)
        # number of step
        nbStep = np.array([4,5,6], dtype=np.int32)
        # degree of the polynomial in each direction
        degree =  np.array([2,1,3], dtype=np.int32)
        # create the Legendre grid
        grid = StOptGrids.RegularLegendreGrid(lowValues,step,nbStep,degree )


        # simulate the perburbed values
        ################################
        nbSimul =40000
        np.random.seed(1000)
        x = np.random.uniform(-1.,1.,size=(1,nbSimul));
        # mesh
        nbMesh = np.array([16],dtype=np.int32)
        # Create the regressor
        #####################
        regressor = StOptReg.LocalLinearRegression(False,x,nbMesh)
        
        # regressed values same values for each point of the grid
        #########################################################
        toReal = (2+x[0,:]+(1+x[0,:])*(1+x[0,:]))
        # function to regress
        toRegress = toReal + 4*np.random.normal(0.,1,nbSimul)
        # create a matrix (number of stock points by number of simulations)
        toRegressMult = np.zeros(shape=(len(toRegress),grid.getNbPoints()))
        for i in range(toRegressMult.shape[1]):
           toRegressMult[:,i] = toRegress        
        # into a list  : two times to test 2 regimes
        listToReg = []
        listToReg.append(toRegressMult)
        listToReg.append(toRegressMult)
    

        # Create the binary archive to dump
        ###################################
        archiveToWrite = StOptGeners.BinaryFileArchive("MyArchive","w")
        # step 1
        archiveToWrite.dumpGridAndRegressedValue("toStore", 1,listToReg, regressor,grid)
        # step 3
        archiveToWrite.dumpGridAndRegressedValue("toStore", 3,listToReg, regressor,grid)


        # Read the regressed values
        ###########################
        archiveToRead =  StOptGeners.BinaryFileArchive("MyArchive","r")
        contValues = archiveToRead.readGridAndRegressedValue(3,"toStore")


        # list of 2 continuation values
        ##############################
        stockPoint = np.array([1.5,3.,5.])
        uncertainty = np.array([0.])
        value =contValues[0].getValue(stockPoint,uncertainty)
      

      # non regular grid
      def testSimpleStorageAndLectureNonRegular(self):
        
        # create the Legendre grid
        grid = StOptGrids.GeneralSpaceGrid([[ 1., 1.7, 2.4, 3.1, 3.8 ],
                                            [2., 4.3, 6.6, 8.9, 11.2, 15.],
                                            [3., 4.9, 5.8, 7.7, 10.,20.]])

        # simulate the perburbed values
        ################################
        nbSimul =40000
        np.random.seed(1000)
        x = np.random.uniform(-1.,1.,size=(1,nbSimul));
        # mesh
        nbMesh = np.array([16],dtype=np.int32)
        # Create the regressor
        #####################
        regressor = StOptReg.LocalLinearRegression(False,x,nbMesh)
        
        # regressed values same values for each point of the grid
        #########################################################
        toReal = (2+x[0,:]+(1+x[0,:])*(1+x[0,:]))
        # function to regress
        toRegress = toReal + 4*np.random.normal(0.,1,nbSimul)
        # create a matrix (number of stock points by number of simulations)
        toRegressMult = np.zeros(shape=(len(toRegress),grid.getNbPoints()))
        for i in range(toRegressMult.shape[1]):
           toRegressMult[:,i] = toRegress        
        # into a list  : two times to test 2 regimes
        listToReg = []
        listToReg.append(toRegressMult)
        listToReg.append(toRegressMult)
    

        # Create the binary archive to dump
        ###################################
        archiveToWrite = StOptGeners.BinaryFileArchive("MyArchive","w")
        # step 1
        archiveToWrite.dumpGridAndRegressedValue("toStore", 1,listToReg, regressor,grid)
        # step 3
        archiveToWrite.dumpGridAndRegressedValue("toStore", 3,listToReg, regressor,grid)


        # Read the regressed values
        ###########################
        archiveToRead =  StOptGeners.BinaryFileArchive("MyArchive","r")
        contValues = archiveToRead.readGridAndRegressedValue(3,"toStore")


        # list of 2 continuation values
        ##############################
        stockPoint = np.array([1.5,3.,5.])
        uncertainty = np.array([0.])
        value =contValues[0].getValue(stockPoint,uncertainty)
        
if __name__ == '__main__': 
    unittest.main()
