#module imports
import numpy
import copy
from fielddisplay import fielddisplay
from EnumDefinitions import *
from StringToEnum import StringToEnum
from checkfield import *
from WriteData import *

class inversion(object):
	"""
	INVERSION class definition

	   Usage:
	      inversion=inversion();
	"""

	#properties
	def __init__(self):
		# {{{ Properties
		self.iscontrol                   = 0
		self.tao                         = 0
		self.incomplete_adjoint          = 0
		self.control_parameters          = float('NaN')
		self.nsteps                      = 0
		self.maxiter_per_step            = float('NaN')
		self.cost_functions              = float('NaN')
		self.cost_functions_coefficients = float('NaN')
		self.gradient_scaling            = float('NaN')
		self.cost_function_threshold     = 0
		self.min_parameters              = float('NaN')
		self.max_parameters              = float('NaN')
		self.step_threshold              = float('NaN')
		self.gradient_only               = 0
		self.vx_obs                      = float('NaN')
		self.vy_obs                      = float('NaN')
		self.vz_obs                      = float('NaN')
		self.vel_obs                     = float('NaN')
		self.thickness_obs               = float('NaN')

		#set defaults
		self.setdefaultparameters()

		#}}}
	def __repr__(self):
		# {{{ Display
		string='\n   Inversion parameters:'
		string="%s\n%s"%(string,fielddisplay(self,'iscontrol','is inversion activated?'))
		string="%s\n%s"%(string,fielddisplay(self,'incomplete_adjoint','do we assume linear viscosity?'))
		string="%s\n%s"%(string,fielddisplay(self,'control_parameters','parameter where inverse control is carried out; ex: {''FrictionCoefficient''}, or {''MaterialsRheologyBbar''}'))
		string="%s\n%s"%(string,fielddisplay(self,'nsteps','number of optimization searches'))
		string="%s\n%s"%(string,fielddisplay(self,'cost_functions','indicate the type of response for each optimization step'))
		string="%s\n%s"%(string,fielddisplay(self,'cost_functions_coefficients','cost_functions_coefficients applied to the misfit of each vertex and for each control_parameter'))
		string="%s\n%s"%(string,fielddisplay(self,'cost_function_threshold','misfit convergence criterion. Default is 1%, NaN if not applied'))
		string="%s\n%s"%(string,fielddisplay(self,'maxiter_per_step','maximum iterations during each optimization step'))
		string="%s\n%s"%(string,fielddisplay(self,'gradient_scaling','scaling factor on gradient direction during optimization, for each optimization step'))
		string="%s\n%s"%(string,fielddisplay(self,'step_threshold','decrease threshold for misfit, default is 30%'))
		string="%s\n%s"%(string,fielddisplay(self,'min_parameters','absolute minimum acceptable value of the inversed parameter on each vertex'))
		string="%s\n%s"%(string,fielddisplay(self,'max_parameters','absolute maximum acceptable value of the inversed parameter on each vertex'))
		string="%s\n%s"%(string,fielddisplay(self,'gradient_only','stop control method solution at gradient'))
		string="%s\n%s"%(string,fielddisplay(self,'vx_obs','observed velocity x component [m/a]'))
		string="%s\n%s"%(string,fielddisplay(self,'vy_obs','observed velocity y component [m/a]'))
		string="%s\n%s"%(string,fielddisplay(self,'vel_obs','observed velocity magnitude [m/a]'))
		string="%s\n%s"%(string,fielddisplay(self,'thickness_obs','observed thickness [m]'))
		string="%s\n%s"%(string,'Available cost functions:')
		string="%s\n%s"%(string,'   101: SurfaceAbsVelMisfit')
		string="%s\n%s"%(string,'   102: SurfaceRelVelMisfit')
		string="%s\n%s"%(string,'   103: SurfaceLogVelMisfit')
		string="%s\n%s"%(string,'   104: SurfaceLogVxVyMisfit')
		string="%s\n%s"%(string,'   105: SurfaceAverageVelMisfit')
		string="%s\n%s"%(string,'   201: ThicknessAbsMisfit')
		string="%s\n%s"%(string,'   501: DragCoefficientAbsGradient')
		string="%s\n%s"%(string,'   502: RheologyBbarAbsGradient')
		string="%s\n%s"%(string,'   503: ThicknessAbsGradient')
		return string
		#}}}

	def setdefaultparameters(self): # {{{
		
		#default is incomplete adjoint for now
		self.incomplete_adjoint=1

		#parameter to be inferred by control methods (only
		#drag and B are supported yet)
		self.control_parameters='FrictionCoefficient'

		#number of steps in the control methods
		self.nsteps=20

		#maximum number of iteration in the optimization algorithm for
		#each step
		self.maxiter_per_step=20*numpy.ones(self.nsteps)

		#the inversed parameter is updated as follows:
		#new_par=old_par + gradient_scaling(n)*C*gradient with C in [0 1];
		#usually the gradient_scaling must be of the order of magnitude of the 
		#inversed parameter (10^8 for B, 50 for drag) and can be decreased
		#after the first iterations
		self.gradient_scaling=50*numpy.ones((self.nsteps,1))

		#several responses can be used:
		self.cost_functions=101*numpy.ones((self.nsteps,1))

		#step_threshold is used to speed up control method. When
		#misfit(1)/misfit(0) < self.step_threshold, we go directly to
		#the next step
		self.step_threshold=.7*numpy.ones(self.nsteps) #30 per cent decrement

		#stop control solution at the gradient computation and return it? 
		self.gradient_only=0

		#cost_function_threshold is a criteria to stop the control methods.
		#if J[n]-J[n-1]/J[n] < criteria, the control run stops
		#NaN if not applied
		self.cost_function_threshold=float('NaN')    #not activated 

		return self
	#}}}

	def checkconsistency(self,md,solution,analyses):    # {{{

		#Early return
		if not self.iscontrol:
			return md

		num_controls=numpy.size(md.inversion.control_parameters)
		num_costfunc=numpy.size(md.inversion.cost_functions,axis=1)

		md = checkfield(md,'inversion.iscontrol','values',[0,1])
		md = checkfield(md,'inversion.tao','values',[0,1])
		md = checkfield(md,'inversion.incomplete_adjoint','values',[0,1])
		md = checkfield(md,'inversion.control_parameters','cell',1,'values',['BalancethicknessThickeningRate','FrictionCoefficient','MaterialsRheologyBbar','MaterialsRheologyZbar','Vx','Vy'])
		md = checkfield(md,'inversion.nsteps','numel',[1],'>=',1)
		md = checkfield(md,'inversion.maxiter_per_step','size',[md.inversion.nsteps],'>=',0)
		md = checkfield(md,'inversion.step_threshold','size',[md.inversion.nsteps])
		md = checkfield(md,'inversion.cost_functions','size',[md.inversion.nsteps,num_costfunc],'values',[101,102,103,104,105,201,501,502,503,504,505])
		md = checkfield(md,'inversion.cost_functions_coefficients','size',[md.mesh.numberofvertices,num_costfunc],'>=',0)
		md = checkfield(md,'inversion.gradient_only','values',[0,1])
		md = checkfield(md,'inversion.gradient_scaling','size',[md.inversion.nsteps,num_controls])
		md = checkfield(md,'inversion.min_parameters','size',[md.mesh.numberofvertices,num_controls])
		md = checkfield(md,'inversion.max_parameters','size',[md.mesh.numberofvertices,num_controls])

		if solution==BalancethicknessSolutionEnum():
			md = checkfield(md,'inversion.thickness_obs','size',[md.mesh.numberofvertices],'NaN',1)
		else:
			md = checkfield(md,'inversion.vx_obs','size',[md.mesh.numberofvertices],'NaN',1)
			md = checkfield(md,'inversion.vy_obs','size',[md.mesh.numberofvertices],'NaN',1)

		return md
	# }}}

	def marshall(self,fid):    # {{{

		WriteData(fid,'object',self,'fieldname','iscontrol','format','Boolean')
		WriteData(fid,'object',self,'fieldname','tao','format','Boolean')
		WriteData(fid,'object',self,'fieldname','incomplete_adjoint','format','Boolean')
		if not self.iscontrol:
			return
		WriteData(fid,'object',self,'fieldname','nsteps','format','Integer')
		WriteData(fid,'object',self,'fieldname','maxiter_per_step','format','DoubleMat','mattype',3)
		WriteData(fid,'object',self,'fieldname','cost_functions_coefficients','format','DoubleMat','mattype',1)
		WriteData(fid,'object',self,'fieldname','gradient_scaling','format','DoubleMat','mattype',3)
		WriteData(fid,'object',self,'fieldname','cost_function_threshold','format','Double')
		WriteData(fid,'object',self,'fieldname','min_parameters','format','DoubleMat','mattype',3)
		WriteData(fid,'object',self,'fieldname','max_parameters','format','DoubleMat','mattype',3)
		WriteData(fid,'object',self,'fieldname','step_threshold','format','DoubleMat','mattype',3)
		WriteData(fid,'object',self,'fieldname','gradient_only','format','Boolean')
		WriteData(fid,'object',self,'fieldname','vx_obs','format','DoubleMat','mattype',1)
		WriteData(fid,'object',self,'fieldname','vy_obs','format','DoubleMat','mattype',1)
		WriteData(fid,'object',self,'fieldname','vz_obs','format','DoubleMat','mattype',1)
		WriteData(fid,'object',self,'fieldname','thickness_obs','format','DoubleMat','mattype',1)

		#process control parameters
		num_control_parameters=len(self.control_parameters)
		data=numpy.array([StringToEnum(control_parameter)[0] for control_parameter in self.control_parameters]).reshape(1,-1)
		WriteData(fid,'data',data,'enum',InversionControlParametersEnum(),'format','DoubleMat','mattype',3)
		WriteData(fid,'data',num_control_parameters,'enum',InversionNumControlParametersEnum(),'format','Integer')

		#process cost functions
		num_cost_functions=numpy.size(self.cost_functions,axis=1)
		data=copy.deepcopy(self.cost_functions)
		data[numpy.nonzero(data==101)]=SurfaceAbsVelMisfitEnum()
		data[numpy.nonzero(data==102)]=SurfaceRelVelMisfitEnum()
		data[numpy.nonzero(data==103)]=SurfaceLogVelMisfitEnum()
		data[numpy.nonzero(data==104)]=SurfaceLogVxVyMisfitEnum()
		data[numpy.nonzero(data==105)]=SurfaceAverageVelMisfitEnum()
		data[numpy.nonzero(data==201)]=ThicknessAbsMisfitEnum()
		data[numpy.nonzero(data==501)]=DragCoefficientAbsGradientEnum()
		data[numpy.nonzero(data==502)]=RheologyBbarAbsGradientEnum()
		data[numpy.nonzero(data==503)]=ThicknessAbsGradientEnum()
		data[numpy.nonzero(data==504)]=ThicknessAlongGradientEnum()
		data[numpy.nonzero(data==505)]=ThicknessAcrossGradientEnum()
		WriteData(fid,'data',data,'enum',InversionCostFunctionsEnum(),'format','DoubleMat','mattype',3)
		WriteData(fid,'data',num_cost_functions,'enum',InversionNumCostFunctionsEnum(),'format','Integer')
	# }}}