# 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 math
import StOptGrids
import StOptReg
import StOptGeners
import matplotlib.pyplot as plt

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

class testBinaryArchiveStOpt(unittest.TestCase):


    def  createData(self):

        X1=np.arange(-1.0 , 1.0 , 0.01 )
        X2=np.arange(-1.0 , 1.0 , 0.01 )
        Y=np.zeros((len(X1),len(X2)))
        for i in range(len(X1)):
            for j in range(len(X2)):
                if i < len(X1)//2:
                    if j < len(X2)//2:
                        Y[i,j]=X1[i]+X2[j]
                    else:
                       Y[i,j]=4*X1[i]+4*X2[j]
                else:
                    if j < len(X2)//2:
                        Y[i,j]=2*X1[i]+1*X2[j]
                    else:
                        Y[i,j]=2*X1[i]+3*X2[j]

        XX1, XX2 = np.meshgrid(X1,X2)  
        Y=Y.T

        r,c = XX1.shape

        X =  np.reshape(XX1,(r*c,1))[:,0]
        I =  np.reshape(XX2,(r*c,1))[:,0]
        Y =  np.reshape(Y,(r*c,1))[:,0]

        xMatrix = np.zeros((2,len(X)))
        xMatrix[0,:] = X
        xMatrix[1,:] = I

        return xMatrix, Y


    def  createData1(self):

        X1=np.arange(-1.0 , 1.0 , 0.01 )
        X2=np.arange(-1.0 , 1.0 , 0.01 )
        Y=np.zeros((len(X1),len(X2)))
        for i in range(len(X1)):
            for j in range(len(X2)):
                if i < len(X1)//2:
                    if j < len(X2)//2:
                        Y[i,j]=2*X1[i]+5*X2[j]
                    else:
                       Y[i,j]=6*X1[i]+6*X2[j]
                else:
                    if j < len(X2)//2:
                        Y[i,j]=0.5*X1[i]+0.5*X2[j]
                    else:
                        Y[i,j]=X1[i]+10*X2[j]

        XX1, XX2 = np.meshgrid(X1,X2)  
        Y=Y.T

        r,c = XX1.shape

        X =  np.reshape(XX1,(r*c,1))[:,0]
        I =  np.reshape(XX2,(r*c,1))[:,0]
        Y =  np.reshape(Y,(r*c,1))[:,0]

        xMatrix = np.zeros((2,len(X)))
        xMatrix[0,:] = X
        xMatrix[1,:] = I

        return xMatrix, Y


    def testSimpleStorageAndLecture(self):
        
        # i create grid of 5 starting from 0 till 4
        lowValues =np.array([0],dtype=float)
        step = np.array([1],dtype=float)
        nbStep = np.array([4], dtype=np.int32)
        grid = StOptGrids.RegularSpaceGrid(lowValues,step,nbStep)

        xMatrix, y = self.createData()

        xMatrix, y1 = self.createData1()

        # create regressor object
        nbMesh = np.array([2,2],dtype=np.int32)
        regressor = StOptReg.LocalLinearRegression(False,xMatrix,nbMesh)
        
        # create a matrix (number of simulations by number of grid points)
        toRegressMult = np.zeros(shape=(len(y),grid.getNbPoints()))

        for i in range(toRegressMult.shape[1]):
            toRegressMult[:,i] = y

        toRegressMult[:,1] = y1


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


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

        # test getMesh1D method
        meshpoints  = regressor.getMesh1D()
                
        
if __name__ == '__main__': 
    unittest.main()
