# 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 pystopt.grids as StOptGrids
import pystopt.regression as StOptReg
import pystopt.geners as 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")

        # print content
        archiveToRead.print(0)
        # detailled
        archiveToRead.print(1)
        
        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()
