# 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

# unit test for regression 
#########################

class testLocalKMeansRegression(unittest.TestCase):

    # One dimensional test
    def testRegression1D(self):
        import StOptReg 
        nbSimul = 5000000;
        np.random.seed(1000)
        x = np.random.uniform(-1.,1.,size=(1,nbSimul));
        # real function
        toReal = (2+x[0,:]+(1+x[0,:])*(1+x[0,:]))
        # function to regress
        toRegress = toReal + 4*np.random.normal(0.,1,nbSimul)
        # mesh
        nbMesh = np.array([20],dtype=np.int32)
        # Regressor
        regressor = StOptReg.LocalLinearRegression(False,x,nbMesh)
        # nb simul
        nbSimul= regressor.getNbSimul()
        # particles
        part = regressor.getParticles()
        # test particules
        y = regressor.getAllSimulations(toRegress).transpose()
        # compare to real value
        diff = 0.5*max(abs((y-toReal)/toReal))
        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")
        

    # two dimensonnal test
    def testRegression2D(self):
        import StOptReg 
        np.random.seed(1000)
        nbSimul = 10000000;
        x = np.random.uniform(-1.,1.,size=(2,nbSimul));
        # real function
        toReal = (2+x[0,:]+(1+x[0,:])*(1+x[0,:]))*(2+x[1,:]+(1+x[1,:])*(1+x[1,:]))
        # function to regress
        toRegress = toReal + 4*np.random.normal(0.,1,nbSimul)
        # mesh
        nbMesh = np.array([16,16],dtype=np.int32)
        # Regressor
        regressor = StOptReg.LocalLinearRegression(False,x,nbMesh)
        y = regressor.getAllSimulations(toRegress).transpose()
        # particles
        part = regressor.getParticles()
        # compare to real value
        diff = max(abs(y-toReal))
        self.assertAlmostEqual(diff,0., 0,"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")

        
    # test different second member
    def testSecondMember(self):
         import StOptReg 
         np.random.seed(1000)
         nbSimul = 100;
         x = np.random.uniform(-1.,1.,size=(1,nbSimul));
         toReal = np.array([((2+x[0,:]+(1+x[0,:])*(1+x[0,:]))).tolist(),((3+x[0,:]+(2+x[0,:])*(2+x[0,:]))).tolist()])
         # function to regress
         toRegress = np.array([(toReal[0,:] + 4*np.random.normal(0.,1,nbSimul)).tolist(),(toReal[1,:] + 4*np.random.normal(0.,1,nbSimul)).tolist()])
         # mesh
         nbMesh = np.array([4],dtype=np.int32)
         # Regressor
         regressor = StOptReg.LocalLinearRegression(False,x,nbMesh)
         y = regressor.getAllSimulationsMultiple(toRegress)
         # check 
         z1 = regressor.getAllSimulations(toRegress[0,:]).transpose()
         z2 = regressor.getAllSimulations(toRegress[1,:]).transpose()
         # compare to real value
         diff = max(abs(y[0,:]-z1[:] ))+max(abs(y[1,:]-z2[:] ))
         self.assertAlmostEqual(diff,0., 7,"Difference between function and its condition expectation estimated greater than tolerance")


    # test  date 0
    def testDateZero(self):
        import StOptReg 
        nbSimul = 50000;
        np.random.seed(1000)
        x = np.random.uniform(-1.,1.,size=(1,nbSimul));
        # real function
        toReal = (2+x[0,:]+(1+x[0,:])*(1+x[0,:]))
        # function to regress
        toRegress = toReal + 4*np.random.normal(0.,1,nbSimul)
        # mesh
        nbMesh = np.array([16],dtype=np.int32)
        # Regressor
        regressor = StOptReg.LocalLinearRegression(True,x,nbMesh)
        y = regressor.getAllSimulations(toRegress).mean()  
        self.assertAlmostEqual(y,toRegress.mean(),8,"Not equality by average")

if __name__ == '__main__': 
    unittest.main()
