# 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 pystopt.grids as StOptGrids
import pystopt.geners as StOptGeners
import pystopt.regression as  StOptReg
import microGrid.simulateState as simX
import microGrid.valueNext as fv


# 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,demand,inventory):

	return fv.finalValue(param,demand,inventory)


def createXMatrix(param,asset, inventory=[]):

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

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



def storageCalculation(param):

	# for regress now 2D, a residual demand process and fixed inventory vector at every time step
	# for gd,  a price process and fixed inventory grid at every time step 
	X0 = np.linspace(-param.ampl,param.ampl,param.nbsimulOpt)
	residualDemand,inventory = simX.residualDemand(param, param.nstep, param.nbsimulOpt, param.maturity, X0, param.rmctype)
	inventory[0] = 0

	if param.rmctype == 'gd':

		stepI = param.I_minMax[1]/param.meshI										
		lowValues =np.array([param.I_minMax[0] , 0.],dtype=np.float64)				
		step = np.array([stepI , 1],dtype=np.float64)  								
		nbStep = np.array([param.meshI , param.H], dtype=np.int32)					
		grid = StOptGrids.RegularSpaceGrid(lowValues,step,nbStep)		

	if param.rmctype == 'regress now 2D':
		
		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)	

	# Value at the final step. 
	valueNext = finalStepValue(param,residualDemand[:,-1],inventory)

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


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

		print("iSteps", iSteps )
		demand = residualDemand[:,iSteps]

		# rmc type and regression method type. The choice is made in the parameter class. 
		if param.rmctype == 'gd':
			nbMesh = np.array([param.meshX],dtype=np.int32)
			regressor = StOptReg.LocalLinearRegression(False,createXMatrix(param,demand),nbMesh)

		elif param.rmctype == 'regress now 2D':
			if param.regType == 'piecewiseLinear':
				nbMesh = np.array([param.meshX, param.meshI],dtype=np.int32)
				regressor = StOptReg.LocalLinearRegression(False,createXMatrix(param,demand, inventory),nbMesh)

			if param.regType == 'globalPolynomial':
				regressor = StOptReg.GlobalCanonicalRegression(False,createXMatrix(param,demand, inventory),param.degree)

			if param.regType == 'kernel':
				regressor = StOptReg.LocalGridKernelRegression(False,createXMatrix(param,demand, inventory),param.bandwidth,param.factPoint, True)

		# Dump the continuation values in the archive:
		archiveToWrite.dumpGridAndRegressedValue("toStore", param.nstep - iSteps, [valueNext], regressor,grid)

		# Read the regressed values
		archiveToRead =  StOptGeners.BinaryFileArchive(param.filetoDump,"r")
		contValues = archiveToRead.readGridAndRegressedValue(param.nstep - iSteps,"toStore")


		valueNext, control_d, StOutput = fv.optimization(contValues[0], demand, inventory, param)


		#######################################################################################################
		# uncomment below to observe following graphs:
		# 1) Cost as a function of residual demand at differnet levels of the inventory.
		# 2) Cost as a function of inventory at differnet levels of the residual demand
		# 3) Control strategy when the generator is ON and when it is OFF.
		# if (iSteps == 0):

		# 	demandCheck = demand #np.tile(demand,param.meshI)
		# 	inventoryCheck = inventory #np.repeat(inventory,param.meshX)

		#  	for q in range(2):
		#  		# plt.figure(q+1)
		#  		plt.scatter(demandCheck,valueNext[:,q],c=inventoryCheck, lw=0, s=10)
		#  		plt.xlabel("demand")
		#  		plt.ylabel("cost")
		#  		name = "q="+str(q)
		#  		plt.title(name)
		#  		plt.colorbar()
		#  		filename = 'q'+str(q)+'_iSteps'+str(iSteps)+'_costDemand.png'
		#  		plt.savefig(filename)
		#  		plt.close()


		#  	for q in range(2):
		#  		# plt.figure(q+1)
		#  		plt.scatter(inventoryCheck,valueNext[:,q],c=demandCheck, lw=0, s=10)
		#  		plt.xlabel("inventory")
		#  		plt.ylabel("cost")
		#  		#plt.ylim(0,1000)
		#  		name = "q="+str(q)
		#  		plt.title(name)
		#  		plt.colorbar()
		#  		filename = 'q'+str(q)+'_iSteps'+str(iSteps)+'_costinventory.png'
		#  		plt.savefig(filename)
		#  		plt.close()

		#  	for q in range(2):
		#  		plt.scatter(demandCheck,inventoryCheck,c=control_d[:,q],lw=0, s=10)
		#  		plt.xlabel('demand')
		#  		plt.ylabel('inventory')
		#  		name = "q="+str(q)
		#  		plt.title(name)
		#  		plt.colorbar()
		#  		filename = 'q'+str(q)+'_iSteps'+str(iSteps)+'_controlMap.png'
		#  		plt.savefig(filename)
		#  		plt.close()

