# Copyright (C) 2021 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


# unit test for regression with some discrete valeus of the state in last dimension
###################################################################################

class testLocalLinearDiscrLastDimRegression(unittest.TestCase):

    # One dimensional test
    # Here two possible state for discrete state
    def testRegression1D(self):
        import StOptReg 
        nbSimul = 500000;
        np.random.seed(1000)
        x = np.zeros(shape=(2,nbSimul))
        # real function
        x[0,:] =  np.random.uniform(-1.,1.,size=(1,nbSimul))
        print(x.dtype)
        stateDiscr =  np.random.randint(2,size=(1,nbSimul))
        x[1,:]  = stateDiscr.astype(float)
        print(x.dtype)
        # real function
        toReal1 = (2+x[0,:]+(1+x[0,:])*(1+x[0,:]))
        toReal2 = (1+x[0,:]+(2+2*x[0,:])*(1+0.1*x[0,:]))
        # function to regress
        toRegress = np.where(stateDiscr[0,:]==0,toReal1, toReal2)
        # mesh (second chosen the number of states  (could be lower))
        nbMesh = np.array([8,2],dtype=np.int32)
        # Regressor
        regressor = StOptReg.LocalLinearDiscrLastDimRegression(False,x,nbMesh)
        # nb simul
        nbSimul= regressor.getNbSimul()
        # particles
        part = regressor.getParticles()
        # test particules
        y = regressor.getAllSimulations(toRegress).transpose()
        # compare to real value
        diff = np.max(np.where(stateDiscr[0,:]==0,np.abs(toReal1-y), np.abs(toReal2-y)))
        self.assertAlmostEqual(diff,0., 1,"Difference between function and its condition expectation estimated greater than tolerance")
        # get back basis function
        regressedFuntionCoeff= regressor.getCoordBasisFunction(toRegress)
        # get back all values
        ySecond= regressor.getValues(x,regressedFuntionCoeff).transpose()
        diff = max(abs(y-ySecond))
        self.assertAlmostEqual(diff,0., 7,"Difference between function and its condition expectation estimated greater than tolerance")
                                        
if __name__ == '__main__': 
    unittest.main()
