import numpy
import copy
from fielddisplay import fielddisplay
from EnumDefinitions import *
from checkfield import checkfield
from WriteData import WriteData
import MatlabFuncs as m

class flowequation(object):
	"""
	FLOWEQUATION class definition

	   Usage:
	      flowequation=flowequation();
	"""

	def __init__(self): # {{{
		
		self.isSIA                          = 0
		self.isSSA                          = 0
		self.isL1L2                         = 0
		self.isHO                           = 0
		self.isFS                           = 0
		self.fe_SSA                         = ''
		self.fe_HO                          = ''
		self.fe_FS                          = ''
		self.augmented_lagrangian_r         = 1.
		self.augmented_lagrangian_rhop      = 1.
		self.augmented_lagrangian_rlambda   = 1.
		self.augmented_lagrangian_rholambda = 1.
		self.XTH_theta                      = 0.
		self.vertex_equation                = float('NaN')
		self.element_equation               = float('NaN')
		self.borderSSA                      = float('NaN')
		self.borderHO                       = float('NaN')
		self.borderFS                       = float('NaN')

		#set defaults
		self.setdefaultparameters()

		#}}}
	def __repr__(self): # {{{
		string='   flow equation parameters:'

		string="%s\n%s"%(string,fielddisplay(self,'isSIA',"is the Shallow Ice Approximation (SIA) used ?"))
		string="%s\n%s"%(string,fielddisplay(self,'isSSA',"is the Shelfy-Stream Approximation (SSA) used ?"))
		string="%s\n%s"%(string,fielddisplay(self,'isL1L2',"are L1L2 equations used ?"))
		string="%s\n%s"%(string,fielddisplay(self,'isHO',"is the Higher-Order (HO) approximation used ?"))
		string="%s\n%s"%(string,fielddisplay(self,'isFS',"are the Full-FS (FS) equations used ?"))
		string="%s\n%s"%(string,fielddisplay(self,'fe_SSA',"Finite Element for SSA: 'P1', 'P1bubble' 'P1bubblecondensed' 'P2'"))
		string="%s\n%s"%(string,fielddisplay(self,'fe_HO' ,"Finite Element for HO:  'P1' 'P1bubble' 'P1bubblecondensed' 'P1xP2' 'P2xP1' 'P2'"))
		string="%s\n%s"%(string,fielddisplay(self,'fe_FS' ,"Finite Element for FS:  'P1P1' (debugging only) 'P1P1GLS' 'MINIcondensed' 'MINI' 'TaylorHood' 'LATaylorHood' 'XTaylorHood'"))
		string="%s\n%s"%(string,fielddisplay(self,'vertex_equation',"flow equation for each vertex"))
		string="%s\n%s"%(string,fielddisplay(self,'element_equation',"flow equation for each element"))
		string="%s\n%s"%(string,fielddisplay(self,'borderSSA',"vertices on SSA's border (for tiling)"))
		string="%s\n%s"%(string,fielddisplay(self,'borderHO',"vertices on HO's border (for tiling)"))
		string="%s\n%s"%(string,fielddisplay(self,'borderFS',"vertices on FS' border (for tiling)"))
		return string
		#}}}
	def setdefaultparameters(self): # {{{

		#P1 for SSA
		self.fe_SSA= 'P1';

		#P1 for HO
		self.fe_HO= 'P1';

		#MINI condensed element for FS by default
		self.fe_FS = 'MINIcondensed';

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

		if StressbalanceAnalysisEnum() in analyses:
			md = checkfield(md,'fieldname','flowequation.isSIA','numel',[1],'values',[0,1])
			md = checkfield(md,'fieldname','flowequation.isSSA','numel',[1],'values',[0,1])
			md = checkfield(md,'fieldname','flowequation.isL1L2','numel',[1],'values',[0,1])
			md = checkfield(md,'fieldname','flowequation.isHO','numel',[1],'values',[0,1])
			md = checkfield(md,'fieldname','flowequation.isFS','numel',[1],'values',[0,1])
			md = checkfield(md,'fieldname','flowequation.fe_SSA','values',['P1','P1bubble','P1bubblecondensed','P2','P2bubble'])
			md = checkfield(md,'fieldname','flowequation.fe_HO' ,'values',['P1','P1bubble','P1bubblecondensed','P1xP2','P2xP1','P2','P2bubble','P1xP3','P2xP4'])
			md = checkfield(md,'fieldname','flowequation.fe_FS' ,'values',['P1P1','P1P1GLS','MINIcondensed','MINI','TaylorHood','XTaylorHood','OneLayerP4z','CrouzeixRaviart'])
			md = checkfield(md,'fieldname','flowequation.borderSSA','size',[md.mesh.numberofvertices],'values',[0,1])
			md = checkfield(md,'fieldname','flowequation.borderHO','size',[md.mesh.numberofvertices],'values',[0,1])
			md = checkfield(md,'fieldname','flowequation.borderFS','size',[md.mesh.numberofvertices],'values',[0,1])
			md = checkfield(md,'fieldname','flowequation.augmented_lagrangian_r','numel',[1],'>',0.)
			md = checkfield(md,'fieldname','flowequation.augmented_lagrangian_rhop','numel',[1],'>',0.)
			md = checkfield(md,'fieldname','flowequation.augmented_lagrangian_rlambda','numel',[1],'>',0.)
			md = checkfield(md,'fieldname','flowequation.augmented_lagrangian_rholambda','numel',[1],'>',0.)
			md = checkfield(md,'fieldname','flowequation.XTH_theta','numel',[1],'>=',0.,'<',.5)
			if m.strcmp(md.mesh.domaintype(),'2Dhorizontal'):
				md = checkfield(md,'fieldname','flowequation.vertex_equation','size',[md.mesh.numberofvertices],'values',[1,2])
				md = checkfield(md,'fieldname','flowequation.element_equation','size',[md.mesh.numberofelements],'values',[1,2])
			elif m.strcmp(md.mesh.domaintype(),'3D'):
				md = checkfield(md,'fieldname','flowequation.vertex_equation','size',[md.mesh.numberofvertices],'values',numpy.arange(0,8+1))
				md = checkfield(md,'fieldname','flowequation.element_equation','size',[md.mesh.numberofelements],'values',numpy.arange(0,8+1))
			else:
				raise RuntimeError('mesh type not supported yet')
			if not (self.isSIA or self.isSSA or self.isL1L2 or self.isHO or self.isFS):
				md.checkmessage("no element types set for this model")

		if StressbalanceSIAAnalysisEnum() in analyses:
			if any(self.element_equation==1):
				if numpy.any(numpy.logical_and(self.vertex_equation,md.mask.groundedice_levelset)):
					print "\n !!! Warning: SIA's model is not consistent on ice shelves !!!\n"

		return md
	# }}}
	def marshall(self,md,fid):    # {{{
		WriteData(fid,'object',self,'fieldname','isSIA','format','Boolean')
		WriteData(fid,'object',self,'fieldname','isSSA','format','Boolean')
		WriteData(fid,'object',self,'fieldname','isL1L2','format','Boolean')
		WriteData(fid,'object',self,'fieldname','isHO','format','Boolean')
		WriteData(fid,'object',self,'fieldname','isFS','format','Boolean')
		WriteData(fid,'enum',FlowequationFeSSAEnum(),'data',StringToEnum(self.fe_SSA)[0],'format','Integer')
		WriteData(fid,'enum',FlowequationFeHOEnum() ,'data',StringToEnum(self.fe_HO)[0] ,'format','Integer')
		WriteData(fid,'enum',FlowequationFeFSEnum() ,'data',StringToEnum(self.fe_FS)[0] ,'format','Integer')
		WriteData(fid,'enum',AugmentedLagrangianREnum(),'data',self.augmented_lagrangian_r ,'format','Double')
		WriteData(fid,'enum',AugmentedLagrangianRhopEnum(),'data',self.augmented_lagrangian_rhop ,'format','Double')
		WriteData(fid,'enum',AugmentedLagrangianRlambdaEnum(),'data',self.augmented_lagrangian_rlambda ,'format','Double')
		WriteData(fid,'enum',AugmentedLagrangianRholambdaEnum(),'data',self.augmented_lagrangian_rholambda ,'format','Double')
		WriteData(fid,'enum',AugmentedLagrangianThetaEnum() ,'data',self.XTH_theta ,'format','Double')
		WriteData(fid,'object',self,'fieldname','borderSSA','format','DoubleMat','mattype',1)
		WriteData(fid,'object',self,'fieldname','borderHO','format','DoubleMat','mattype',1)
		WriteData(fid,'object',self,'fieldname','borderFS','format','DoubleMat','mattype',1)
		#convert approximations to enums
		data=copy.deepcopy(self.vertex_equation)
		data[numpy.nonzero(data==0)]=NoneApproximationEnum()
		data[numpy.nonzero(data==1)]=SIAApproximationEnum()
		data[numpy.nonzero(data==2)]=SSAApproximationEnum()
		data[numpy.nonzero(data==3)]=L1L2ApproximationEnum()
		data[numpy.nonzero(data==4)]=HOApproximationEnum()
		data[numpy.nonzero(data==5)]=FSApproximationEnum()
		data[numpy.nonzero(data==6)]=SSAHOApproximationEnum()
		data[numpy.nonzero(data==7)]=HOFSApproximationEnum()
		data[numpy.nonzero(data==8)]=SSAFSApproximationEnum()
		WriteData(fid,'data',data,'enum',FlowequationVertexEquationEnum(),'format','DoubleMat','mattype',1)
		data=copy.deepcopy(self.element_equation)
		data[numpy.nonzero(data==0)]=NoneApproximationEnum()
		data[numpy.nonzero(data==1)]=SIAApproximationEnum()
		data[numpy.nonzero(data==2)]=SSAApproximationEnum()
		data[numpy.nonzero(data==3)]=L1L2ApproximationEnum()
		data[numpy.nonzero(data==4)]=HOApproximationEnum()
		data[numpy.nonzero(data==5)]=FSApproximationEnum()
		data[numpy.nonzero(data==6)]=SSAHOApproximationEnum()
		data[numpy.nonzero(data==7)]=SSAFSApproximationEnum()
		data[numpy.nonzero(data==8)]=HOFSApproximationEnum()
		WriteData(fid,'data',data,'enum',FlowequationElementEquationEnum(),'format','DoubleMat','mattype',2)
	# }}}
