import numpy
import copy
from fielddisplay import fielddisplay
from EnumDefinitions import *
from StringToEnum import StringToEnum
from checkfield import checkfield
from WriteData import WriteData

class m1qn3inversion(object):
	'''
	M1QN3 class definition

   Usage:
      m1qn3inversion=m1qn3inversion()
	'''

	def __init__(self,*args): # {{{

		if not len(args):
			print 'empty init'
			self.iscontrol                   = 0
			self.incomplete_adjoint          = 0
			self.control_parameters          = float('NaN')
			self.maxsteps                    = 0
			self.maxiter                     = 0
			self.dxmin                       = 0.
			self.gttol                       = 0.
			self.cost_functions              = float('NaN')
			self.cost_functions_coefficients = float('NaN')
			self.min_parameters              = float('NaN')
			self.max_parameters              = float('NaN')
			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()
		elif len(args)==1 and args[0].__module__=='inversion':
			print 'converting inversion to m1qn3inversion'
			inv=args[0]
			self.iscontrol                   = inv.iscontrol
			self.incomplete_adjoint          = inv.incomplete_adjoint
			self.control_parameters          = inv.control_parameters
			self.maxsteps                    = inv.nsteps
			self.cost_functions              = inv.cost_functions
			self.cost_functions_coefficients = inv.cost_functions_coefficients
			self.min_parameters              = inv.min_parameters
			self.max_parameters              = inv.max_parameters
			self.vx_obs                      = inv.vx_obs
			self.vy_obs                      = inv.vy_obs
			self.vz_obs                      = inv.vz_obs
			self.vel_obs                     = inv.vel_obs
			self.thickness_obs               = inv.thickness_obs
		else:
			raise Exception('constructor not supported')
		#}}}
	def __repr__(self): # {{{
		string='   m1qn3inversion parameters:'
		string="%s\n%s"%(string,fielddisplay(self,'iscontrol','is inversion activated?'))
		string="%s\n%s"%(string,fielddisplay(self,'incomplete_adjoint','1: linear viscosity, 0: non-linear viscosity'))
		string="%s\n%s"%(string,fielddisplay(self,'control_parameters','ex: [''FrictionCoefficient''], or [''MaterialsRheologyBbar'']'))
		string="%s\n%s"%(string,fielddisplay(self,'maxsteps','maximum number of iterations (gradient computation)'))
		string="%s\n%s"%(string,fielddisplay(self,'maxiter','maximum number of Function evaluation (forward run)'))
		string="%s\n%s"%(string,fielddisplay(self,'dxmin','convergence criterion: two points less than dxmin from eachother (sup-norm) are considered identical'))
		string="%s\n%s"%(string,fielddisplay(self,'gttol','||g(X)||/||g(X0)|| (g(X0): gradient at initial guess X0)'))
		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,'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,'vx_obs','observed velocity x component [m/yr]'))
		string="%s\n%s"%(string,fielddisplay(self,'vy_obs','observed velocity y component [m/yr]'))
		string="%s\n%s"%(string,fielddisplay(self,'vel_obs','observed velocity magnitude [m/yr]'))
		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 iterations
		self.maxsteps=20
		self.maxiter=40

		#several responses can be used:
		self.cost_functions=101

		#m1qn3 parameters
		self.dxmin  = 0.1
		self.gttol = 1e-4

		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)

		md = checkfield(md,'fieldname','inversion.iscontrol','values',[0,1])
		md = checkfield(md,'fieldname','inversion.incomplete_adjoint','values',[0,1])
		md = checkfield(md,'fieldname','inversion.control_parameters','cell',1,'values',['BalancethicknessThickeningRate','FrictionCoefficient','MaterialsRheologyBbar','DamageDbar','Vx','Vy','Thickness'])
		md = checkfield(md,'fieldname','inversion.maxsteps','numel',[1],'>=',0)
		md = checkfield(md,'fieldname','inversion.maxiter','numel',[1],'>=',0)
		md = checkfield(md,'fieldname','inversion.dxmin','numel',[1],'>',0.)
		md = checkfield(md,'fieldname','inversion.gttol','numel',[1],'>',0.)
		md = checkfield(md,'fieldname','inversion.cost_functions','size',[num_costfunc],'values',[101,102,103,104,105,201,501,502,503,504,505])
		md = checkfield(md,'fieldname','inversion.cost_functions_coefficients','size',[md.mesh.numberofvertices,num_costfunc],'>=',0)
		md = checkfield(md,'fieldname','inversion.min_parameters','size',[md.mesh.numberofvertices,num_controls])
		md = checkfield(md,'fieldname','inversion.max_parameters','size',[md.mesh.numberofvertices,num_controls])

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

		return md
	# }}}
	def marshall(self,md,fid):    # {{{

		yts=365.0*24.0*3600.0

		WriteData(fid,'object',self,'class','inversion','fieldname','iscontrol','format','Boolean')
		WriteData(fid,'enum',InversionTypeEnum(),'data',2,'format','Integer')
		if not self.iscontrol:
			return
		WriteData(fid,'object',self,'class','inversion','fieldname','incomplete_adjoint','format','Boolean')
		WriteData(fid,'object',self,'class','inversion','fieldname','maxsteps','format','Integer')
		WriteData(fid,'object',self,'class','inversion','fieldname','maxiter','format','Integer')
		WriteData(fid,'object',self,'class','inversion','fieldname','dxmin','format','Double')
		WriteData(fid,'object',self,'class','inversion','fieldname','gttol','format','Double')
		WriteData(fid,'object',self,'class','inversion','fieldname','cost_functions_coefficients','format','DoubleMat','mattype',1)
		WriteData(fid,'object',self,'class','inversion','fieldname','min_parameters','format','DoubleMat','mattype',3)
		WriteData(fid,'object',self,'class','inversion','fieldname','max_parameters','format','DoubleMat','mattype',3)
		WriteData(fid,'object',self,'class','inversion','fieldname','vx_obs','format','DoubleMat','mattype',1,'scale',1./yts)
		WriteData(fid,'object',self,'class','inversion','fieldname','vy_obs','format','DoubleMat','mattype',1,'scale',1./yts)
		WriteData(fid,'object',self,'class','inversion','fieldname','vz_obs','format','DoubleMat','mattype',1,'scale',1./yts)
		WriteData(fid,'object',self,'class','inversion','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)
		data=copy.deepcopy(self.cost_functions)
		pos=[i for i,x in enumerate(self.cost_functions) if x==101];
		for i in pos: data[i]=SurfaceAbsVelMisfitEnum()        
		pos=[i for i,x in enumerate(self.cost_functions) if x==102];
		for i in pos: data[i]=SurfaceRelVelMisfitEnum()        
		pos=[i for i,x in enumerate(self.cost_functions) if x==103];
		for i in pos: data[i]=SurfaceLogVelMisfitEnum()        
		pos=[i for i,x in enumerate(self.cost_functions) if x==104];
		for i in pos: data[i]=SurfaceLogVxVyMisfitEnum()       
		pos=[i for i,x in enumerate(self.cost_functions) if x==105];
		for i in pos: data[i]=SurfaceAverageVelMisfitEnum()    
		pos=[i for i,x in enumerate(self.cost_functions) if x==201];
		for i in pos: data[i]=ThicknessAbsMisfitEnum()         
		pos=[i for i,x in enumerate(self.cost_functions) if x==501];
		for i in pos: data[i]=DragCoefficientAbsGradientEnum() 
		pos=[i for i,x in enumerate(self.cost_functions) if x==502];
		for i in pos: data[i]=RheologyBbarAbsGradientEnum()    
		pos=[i for i,x in enumerate(self.cost_functions) if x==503];
		for i in pos: data[i]=ThicknessAbsGradientEnum()       
		pos=[i for i,x in enumerate(self.cost_functions) if x==504];
		for i in pos: data[i]=ThicknessAlongGradientEnum()     
		pos=[i for i,x in enumerate(self.cost_functions) if x==505];
		for i in pos: data[i]=ThicknessAcrossGradientEnum()    
		WriteData(fid,'data',numpy.array(data).reshape(1,-1),'enum',InversionCostFunctionsEnum(),'format','DoubleMat','mattype',3)
		WriteData(fid,'data',num_cost_functions,'enum',InversionNumCostFunctionsEnum(),'format','Integer')
	# }}}
