import numpy as npy
from pairoptions import pairoptions
from averaging import averaging
from plotmodel import plotmodel

def analyticaldamage(md,*args):
	'''
	ANALYTICALDAMAGE - compute damage for an ice shelf 
	
		 This routine computes damage as a function of water/ice
		 material properties, ice thickness, strain rate, and ice 
		 rigidity.  The model must contain computed strain rates,
		 either from observed or modeled ice velocities.
	
	   Available options:
			- 'eq'			: analytical equation to use in the calculation.  Must be one of:
									'Weertman1D' for a confined ice shelf free to flow in one direction
									'Weertman2D' for an unconfined ice shelf free to spread in any direction
									'Thomas' for a 2D ice shelf, taking into account full strain rate tensor (default)
			- 'smoothing'	: the amount of smoothing to be applied to the strain rate data.
									Type 'help averaging' for more information on its usage.
			- 'sigmab'		: a compressive backstress term to be subtracted from the driving stress 
									in the damage calculation
	
	   Return values:
			'damage' which is truncated in the range [0,1-1e-9]
	
		   'B' is the rigidity, which is equal to md.materials.rheology_B in areas outside
			those defined by 'mask.'  Within areas defined by 'mask,' where negative damage 
			is inferred, 'B' is updated to make damage equal to zero.  
	
			'backstress' is the inferred backstress necessary to balance the analytical solution
			(keeping damage within its appropriate limits, e.g. D in [0,1]).
	
	   Usage:
	      [damage,B,backstress]=analyticaldamage(md,options)
	
	   Example:
	      [damage,B,backstress]=analyticaldamage(md,'eq','Weertman2D','smoothing',2,'backstress',10e3)
	'''

	# check inputs
	if 'strainrate' not in md.results.__dict__:
		raise StandardError('md.results.strainrate not present.  Calculate using md=mechanicalproperties(md,vx,vy)')
	if md.mesh.dimension!=2:
		raise StandardError('only 2D model supported currently')
	if npy.any(md.flowequation.element_equation!=2):
		print 'Warning: the model has some non SSA elements. These will be treated like SSA elements'

	# process options
	options = pairoptions(*args)
	eq = options.getfieldvalue('eq','Thomas')
	smoothing = options.getfieldvalue('smoothing',0)
	sigmab = options.getfieldvalue('sigmab',0)
	if len(sigmab==1):
		sigmab=sigmab*npy.ones(md.mesh.numberofvertices,)
	
	# average element strain rates onto vertices
	e1=averaging(md,md.results.strainrate.principalvalue1,smoothing)/md.constants.yts # convert to s^-1
	e2=averaging(md,md.results.strainrate.principalvalue2,smoothing)/md.constants.yts
	exx=averaging(md,md.results.strainrate.xx,smoothing)/md.constants.yts
	eyy=averaging(md,md.results.strainrate.yy,smoothing)/md.constants.yts
	exy=averaging(md,md.results.strainrate.xy,smoothing)/md.constants.yts
	
	# checks: any of e1 or e2 equal to zero?
	pos=npy.nonzero(e1==0)
	if npy.any(pos==1):
		print 'WARNING: first principal strain rate equal to zero.  Value set to 1e-13 s^-1'
		e1[pos]=1e-13
	pos=npy.nonzero(e2==0)
	if npy.any(pos==1):
		disp('WARNING: second principal strain rate equal to zero.  Value set to 1e-13 s^-1');
		e2[pos]=1e-13
	
	## old method using principal strain rates {{{
	## ex=maximum principal tensile strain rate
	#ex=e1;
	#a=e2./e1;
	#pos=find(e1<0 & e2>0); # longitudinal compression and lateral tension
	#a(pos)=e1(pos)./e2(pos);
	#ex(pos)=e2(pos);
	#pos2=find(e1<0 & e2<0 & abs(e1)<abs(e2)); # lateral and longitudinal compression
	#a(pos2)=e1(pos2)./e2(pos2);
	#ex(pos2)=e2(pos2);
	#pos3=find(e1>0 & e2>0 & abs(e1)<abs(e2)); # lateral and longitudinal tension 
	#a(pos3)=e1(pos3)./e2(pos3);
	#ex(pos3)=e2(pos3);
	#id=find(e1<0 & e2<0);
	#a(id)=-a(id); # where both strain rates are compressive, enforce negative alpha
	#
	## }}}
	
	# new method using longitudinal strain rates defined by observed velocity vector
	velangle=npy.arctan(md.initialization.vy/md.initialization.vx)
	ex=0.5*(exx+eyy)+0.5*(exx-eyy)*npy.cos(2.*velangle)+exy*npy.sin(2.*velangle)
	ey=exx+eyy-ex # trace of strain rate tensor is invariant
	exy=-0.5*(exx-eyy)*npy.sin(2.*velangle)+exy*npy.cos(2.*velangle)
	a=ey/ex
	b=exy/ex
	pos=npy.nonzero(npy.logical_and(ex<0,ey<0))
	#length(pos)
	a[pos]=-a[pos]
	
	# a < -1 in areas of strong lateral compression or longitudinal compression
	# and theta is undefined at a = -2
	pos=npy.nonzero(abs((abs(a)-2))<1e-3)
	a[pos]=-2+1e-3
	
	# spreading stress
	rhoi=md.materials.rho_ice
	rhow=md.materials.rho_water
	C=0.5*rhoi*md.constants.g*(1.-rhoi/rhow)
	T=C*md.geometry.thickness
	
	# rheology
	B=md.materials.rheology_B
	n=averaging(md,md.materials.rheology_n,0)
	
	if eq=='Weertman1D':
		theta=1./8*npy.ones(md.mesh.numberofvertices,)
		a=npy.zeros(md.mesh.numberofvertices,)
	elif eq=='Weertman2D':
		theta=1./9*npy.ones(md.mesh.numberofvertices,1)
		a=npy.ones(md.mesh.numberofvertices,)
	elif eq=='Thomas':
		theta=((1+a+a**2+b**2)**((n-1)/2))/(abs(2+a)**n)
	else:
		raise StandardError('argument passed to "eq" not valid.  Type "help analyticaldamage" for usage')
	
	D=1-(1+a+a**2+b**2)**((n-1)/(2*n))/abs(ex)**(1./n)*(T-sigmab)/B/(2+a)/npy.sign(ex)
	
	backstress=npy.zeros(md.mesh.numberofvertices,)
	
	# D>1 where (2+a).*sign(ex)<0, compressive regions where high backstress needed
	pos=npy.nonzero(D>1)
	D[pos]=0
	
	# backstress to bring damage to zero
	backstress[pos]=T[pos]-(1-D[pos])*B[pos]*npy.sign(ex[pos])*(2+a[pos])*abs(ex[pos])**(1./n[pos])/(1+a[pos]+a[pos]**2)**((n[pos]-1)/2./n[pos])
	
	pos=npy.nonzero(D<0)
	D[pos]=0

	# backstress to bring negative damage to zero
	backstress[pos]=T[pos]-(1-D[pos])*B[pos]*npy.sign(ex[pos])*(2+a[pos])*abs(ex[pos])**(1./n[pos])/(1+a[pos]+a[pos]**2)**((n[pos]-1)/2./n[pos])
	
	pos=npy.nonzero(backstress<0)
	backstress[pos]=0
	
	# increased rigidity to bring negative damage to zero
	B[pos]=npy.sign(ex[pos])/(2+a[pos])*(1+a[pos]+a[pos]**2)**((n[pos]-1)/2./n[pos])*T[pos]/abs(ex[pos])**(1./n[pos]);
	
	damage=D
	
	return damage, B, backstress
