# 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

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

class testSparseRegression(unittest.TestCase):

    # One dimensional test
    def testSparseRegression1D(self):
        import StOptReg
        nbSimul = 2000000;
        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)
        # level for sparse grid
        iLevel = 5;
        # weight for anisotropic sparse grids
        weight= np.array([1],dtype=np.int32)
        # Regressor degree 1
        regressor = StOptReg.SparseRegression(False,x,iLevel, weight, 1)
        y = regressor.getAllSimulations(toRegress).transpose()
        # compare to real value
        diff = max(abs(y-toReal))
        self.assertAlmostEqual(diff,0., 1,"Difference between function and its condition expectation estimated greater than tolerance")
        # Regressor degree 2
        regressor = StOptReg.SparseRegression(False,x,iLevel, weight, 2)
        y = regressor.getAllSimulations(toRegress).transpose()
        # compare to real value
        diff = max(abs(y-toReal))
        self.assertAlmostEqual(diff,0., 1,"Difference between function and its condition expectation estimated greater than tolerance")
        # Regressor degree 3
        regressor = StOptReg.SparseRegression(False,x,iLevel, weight, 3)
        y = regressor.getAllSimulations(toRegress).transpose()
        # compare to real value
        diff = max(abs(y-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 testSparseRegression2D(self):
        import StOptReg
        np.random.seed(1000)
        nbSimul = 1000000;
        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)
        # level for sparse grid
        iLevel = 5;
        # weight for anisotropic sparse grids
        weight= np.array([1,1],dtype=np.int32)
        # Regressor degree 1
        regressor = StOptReg.SparseRegression(False,x,iLevel, weight, 1)
        y = regressor.getAllSimulations(toRegress).transpose()
        # 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")
        # Regressor degree 2
        regressor = StOptReg.SparseRegression(False,x,iLevel, weight, 2)
        y = regressor.getAllSimulations(toRegress).transpose()
        # 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")
        # Regressor degree 3
        regressor = StOptReg.SparseRegression(False,x,iLevel, weight, 2)
        y = regressor.getAllSimulations(toRegress).transpose()
        # 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)
        # level for sparse grid
        iLevel = 5;
        # weight for anisotropic sparse grids
        weight= np.array([1],dtype=np.int32)
        # Regressor degree 1
        regressor = StOptReg.SparseRegression(True,x,iLevel, weight, 1)
        y = regressor.getAllSimulations(toRegress).transpose()
        self.assertAlmostEqual(y[0],toRegress.mean(),8,"Not equality by average")

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