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

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.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'"))
		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,'flowequation.isSIA','numel',[1],'values',[0,1])
			md = checkfield(md,'flowequation.isSSA','numel',[1],'values',[0,1])
			md = checkfield(md,'flowequation.isL1L2','numel',[1],'values',[0,1])
			md = checkfield(md,'flowequation.isHO','numel',[1],'values',[0,1])
			md = checkfield(md,'flowequation.isFS','numel',[1],'values',[0,1])
			md = checkfield(md,'flowequation.fe_SSA','values',['P1','P1bubble','P1bubblecondensed','P2'])
			md = checkfield(md,'flowequation.fe_HO' ,'values',['P1','P1bubble','P1bubblecondensed','P1xP2','P2xP1','P2'])
			md = checkfield(md,'flowequation.fe_FS' ,'values',['P1P1','P1P1GLS','MINIcondensed','MINI','TaylorHood'])
			md = checkfield(md,'flowequation.borderSSA','size',[md.mesh.numberofvertices],'values',[0,1])
			md = checkfield(md,'flowequation.borderHO','size',[md.mesh.numberofvertices],'values',[0,1])
			md = checkfield(md,'flowequation.borderFS','size',[md.mesh.numberofvertices],'values',[0,1])
			if md.mesh.dimension==2:
				md = checkfield(md,'flowequation.vertex_equation','size',[md.mesh.numberofvertices],'values',[1,2])
				md = checkfield(md,'flowequation.element_equation','size',[md.mesh.numberofelements],'values',[1,2])
			else:
				md = checkfield(md,'flowequation.vertex_equation','size',[md.mesh.numberofvertices],'values',numpy.arange(0,8+1))
				md = checkfield(md,'flowequation.element_equation','size',[md.mesh.numberofelements],'values',numpy.arange(0,8+1))
			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,'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)]=HOApproximationEnum()
		data[numpy.nonzero(data==4)]=FSApproximationEnum()
		data[numpy.nonzero(data==5)]=SSAHOApproximationEnum()
		data[numpy.nonzero(data==6)]=SSAFSApproximationEnum()
		data[numpy.nonzero(data==7)]=HOFSApproximationEnum()
		data[numpy.nonzero(data==8)]=L1L2ApproximationEnum()
		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)]=HOApproximationEnum()
		data[numpy.nonzero(data==4)]=FSApproximationEnum()
		data[numpy.nonzero(data==5)]=SSAHOApproximationEnum()
		data[numpy.nonzero(data==6)]=SSAFSApproximationEnum()
		data[numpy.nonzero(data==7)]=HOFSApproximationEnum()
		data[numpy.nonzero(data==8)]=L1L2ApproximationEnum()
		WriteData(fid,'data',data,'enum',FlowequationElementEquationEnum(),'format','DoubleMat','mattype',2)
	# }}}
