# Copyright (C) 2017 EDF
# All Rights Reserved
# This code is published under the GNU Lesser General Public License (GNU LGPL)
import numpy as np
import StOptReg 
import StOptGeners

# unit test to show how to store some regression object and basis function coefficients associated
##################################################################################################

def  createData():

	X1=np.arange(0.0 , 2.2 , 0.01 )
	X2=np.arange(0.0 , 1.1 , 0.005 )
	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]+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




# main

xMatrix, y = createData()

# 2 dimensional regression 2 by 2 meshes
nbMesh = np.array([2,2],dtype=np.int32)
regressor = StOptReg.LocalLinearRegression(False,xMatrix,nbMesh)

# coefficients
coeff = regressor.getCoordBasisFunction(y)

print("Regressed coeff", coeff) 


# store them in a matrix
coeffList = np.zeros(shape=(1,3*2*2))
coeffList[0,:]=coeff.transpose()


# archive write for  regressors
archiveWriteForRegressor =StOptGeners.BinaryFileArchive("archive","w")

# store
step =1
archiveWriteForRegressor.dumpSome2DArray("RegCoeff",step,coeff)
archiveWriteForRegressor.dumpSomeRegressor("Regressor",step,regressor)
#  store a one D array
name1 = "AllSteps"
name2 = "Product"
for i in range(5):
        x = i*np.ones(10)
        archiveWriteForRegressor.dumpSome1DArray(name1, name2, x)

# archive Read for  regressors
archiveReadForRegressor =StOptGeners.BinaryFileArchive("archive","r")

# get back
values = archiveReadForRegressor.readSome2DArray("RegCoeff",step)
reg = archiveReadForRegressor.readSomeRegressor("Regressor",step)
print("Regressed coeff ", values)
print("Reg",reg)

for i in range(5):
        x =archiveReadForRegressor.readSome1DArray(name1, name2, i)
        print("X ", x)


