# Copyright (C) 2018 The Regents of the University of California, Michael Ludkovski and Aditya Maheshwari
# All Rights Reserved
# This code is published under the GNU Lesser General Public License (GNU LGPL)

from __future__ import division
import numpy as np
import math
import matplotlib.pyplot as plt
import time

import StOptGrids
import StOptReg
import StOptGeners

import microgridDEA.simulateState_new as simX
import microgridDEA.valueNext_new as fv
import microgridDEA.sobol_lib as sl


# returns the value function at the final step. Depending upon the type of RMC method, the size of the matrix may be different.
def finalStepValue(param,regParams,demand,inventory):

	return fv.finalValue(param,regParams,demand,inventory)

# creates the matrix using residual demand and inventory for the StOpt function.
def createXMatrix(param,regParams,asset, inventory=[]):

	if regParams.rmctype == 'regress now 2D':
		xMatrix = np.zeros((2,len(asset)))
		xMatrix[0,:] = asset
		xMatrix[1,:] = inventory
		return xMatrix

	if regParams.rmctype == 'gd':
		xMatrix = np.zeros((1,len(asset)))
		xMatrix[0,:] = asset
		return xMatrix

def gridType(param,regParams):

	if regParams.rmctype == 'regress now 2D':
		# notice the grid is only for the regime
		lowValues =np.array([0.],dtype=np.float64)				            # low value for the mesh
		step = np.array([1],dtype=np.float64)  								# size of the mesh
		nbStep = np.array([param.H], dtype=np.int32)						# number of step
		grid = StOptGrids.RegularSpaceGrid(lowValues,step,nbStep)	

	if regParams.rmctype == 'gd':
		# notice the grid is for both the regime and the inventory levels.
		stepI = param.I_minMax[1]/regParams.meshI										
		lowValues =np.array([param.I_minMax[0] , 0.],dtype=np.float64)				
		step = np.array([stepI , 1],dtype=np.float64)  								
		nbStep = np.array([regParams.meshI , param.H], dtype=np.int32)					
		grid = StOptGrids.RegularSpaceGrid(lowValues,step,nbStep)		

	return grid

def regObject(param,regParams,x):

	# rmc type and regression type. The choice is made in the "basicVariables" and "regressionType" class. 
	if regParams.rmctype == 'gd':
		nbMesh = np.array([regParams.meshX],dtype=np.int32)
		regressor = StOptReg.LocalLinearRegression(False,x,nbMesh)

	elif regParams.rmctype == 'regress now 2D':
		if regParams.regType == 'globalPolynomial':
			regressor = StOptReg.GlobalCanonicalRegression(False,x,regParams.degree)

	return regressor


def designSites(param,regParams):

	if regParams.rmctype == 'regress now 2D':
		dim = 2
		# number of unique design sites.
		numdesignSites = regParams.nbUniqueSite 
		skip=0 
		# use sobol sequence for the design sites. This could be replaced by LHS design. 
		sites = sl.i4_sobol_generate(dim, numdesignSites, skip)

		# sobol sequences are generated on square domain. extend it to the domain of interest.
		X0 = -param.ampl + 2*param.ampl*sites[0,:]
		I0 = param.I_minMax[0] + (param.I_minMax[1]-param.I_minMax[0])*sites[1,:]
		
		# repeat the design sites depending upon the batch size.
		X0_rep = np.repeat(X0,regParams.batchSize)
		I0_rep = np.repeat(I0,regParams.batchSize)

		return X0,I0, X0_rep, I0_rep

	elif regParams.rmctype == 'gd':
		dim = 1 # only one dimension because of rmc type.

		numdesignSites = regParams.nbUniqueSite # number of unique design sites
		skip=0
		sites = sl.i4_sobol_generate(dim, numdesignSites, skip ) 

		# sobol sequence in demand dimension and uniformly spaced in the inventory. 		
		X0 = -param.ampl + 2*param.ampl*sites[0,:]
		I0 = np.linspace(param.I_minMax[0],param.I_minMax[1],regParams.meshI+1)

		# repeat the design sites depending upon the batch size.
		X0_rep = np.repeat(X0,regParams.batchSize)

		return X0,I0,X0_rep,I0


def storageCalculation(param, regParamObjcs):

	# parameters of the regression for example: local linear or global polynomial and type of grid.
	regParams = regParamObjcs[-1]

	# create design sites
	X0, inventory, X0_rep, inventory_rep = designSites(param,regParams) 

	# generate one step residual demand
	residualDemand = simX.residualDemand(param, 1, regParams.nbsimulOpt, param.dt, X0_rep, regParams.rmctype)

	# compute the value at the final step. 
	valueNext = finalStepValue(param,regParams,residualDemand[:,1],inventory)

	# average the value for each batch
	r,c = valueNext.shape
	valueNext = np.mean(np.reshape(valueNext, (regParams.batchSize,	regParams.nbUniqueSite, c), order='F'),0)

	archiveToWrite = StOptGeners.BinaryFileArchive(param.filetoDump,"w")


	# iterate on time steps
	for iSteps in range(param.nstep-1,-1,-1):
		print("iSteps",iSteps)

		# create grid for the regression
		grid = gridType(param,regParamObjcs[iSteps])
		
		# create regressor object
		regressor = regObject(param,regParamObjcs[iSteps],createXMatrix(param,regParams,X0,inventory))
		
		# Dump the continuation values in the archive:
		archiveToWrite.dumpGridAndRegressedValue("toStore", param.nstep - iSteps, [valueNext], regressor,grid)

		if iSteps>0:
			regParams = regParamObjcs[iSteps-1]

			# create design sites
			X0, inventory, X0_rep, inventory_rep = designSites(param,regParams) # isites is useless if we use rmctype = 'gd'.

			# generate one step residual demand
			residualDemand = simX.residualDemand(param, 1, regParams.nbsimulOpt, param.dt, X0_rep, regParams.rmctype)
			
			# Read the regressed values
			archiveToRead =  StOptGeners.BinaryFileArchive(param.filetoDump,"r")
			contValues = archiveToRead.readGridAndRegressedValue(param.nstep - iSteps,"toStore")

			# regParamObjcs[iSteps] tells the function how the regression was done and regParamObjcs[iSteps-1] tell the function how to store the valueNext now.
			valueNext, control_d, StOutput = fv.optimization(contValues[0], residualDemand[:,1], inventory_rep, param, regParamObjcs[iSteps-1], regParamObjcs[iSteps])
			
			# average the value for each batch
			r,c = valueNext.shape
			valueNext = np.mean(np.reshape(valueNext, (regParams.batchSize,	regParams.nbUniqueSite, c), order='F'),0)