import os
import numpy
from ContourToMesh import *

def SetMarineIceSheetBC(md,icefrontfile=''):
	"""
	SETICEMARINESHEETBC - Create the boundary conditions for diagnostic and thermal models for a  Marine Ice Sheet with Ice Front

	   Neumann BC are used on the ice front (an ARGUS contour around the ice front
	   can be given in input, or it will be deduced as onfloatingice & onboundary)
	   Dirichlet BC are used elsewhere for diagnostic

	   Usage:
	      md=SetMarineIceSheetBC(md,icefrontfile)
	      md=SetMarineIceSheetBC(md)

	   Example:
	      md=SetMarineIceSheetBC(md,'Front.exp')
	      md=SetMarineIceSheetBC(md)

	   See also: SETICESHELFBC, SETMARINEICESHEETBC
	"""

	#node on Dirichlet (boundary and ~icefront)
	if icefrontfile:
		#User provided Front.exp, use it
		if not os.path.exists(icefrontfile):
			raise IOError("SetMarineIceSheetBC error message: ice front file '%s' not found." % icefrontfile)
		[nodeinsideicefront,dum]=ContourToMesh(md.mesh.elements,md.mesh.x.reshape(-1,1),md.mesh.y.reshape(-1,1),icefrontfile,'node',2)
		vertexonicefront=numpy.logical_and(md.mesh.vertexonboundary,nodeinsideicefront.reshape(-1)).astype(float)
	else:
		#Guess where the ice front is
		vertexonfloatingice=numpy.zeros((md.mesh.numberofvertices))
		vertexonfloatingice[md.mesh.elements[numpy.nonzero(md.mask.elementonfloatingice),:].astype(int)-1]=1
		vertexonicefront=numpy.logical_and(md.mesh.vertexonboundary,vertexonfloatingice).astype(float)

#	pos=find(md.mesh.vertexonboundary & ~vertexonicefront);
	pos=numpy.nonzero(numpy.logical_and(md.mesh.vertexonboundary,numpy.logical_not(vertexonicefront)))[0]
	if not numpy.size(pos):
		print "SetMarineIceSheetBC warning: ice front all around the glacier, no dirichlet found. Dirichlet must be added manually."

	md.diagnostic.spcvx=float('nan')*numpy.ones((md.mesh.numberofvertices,1))
	md.diagnostic.spcvy=float('nan')*numpy.ones((md.mesh.numberofvertices,1))
	md.diagnostic.spcvz=float('nan')*numpy.ones((md.mesh.numberofvertices,1))
	md.diagnostic.spcvx[pos]=0
	md.diagnostic.spcvy[pos]=0
	md.diagnostic.spcvz[pos]=0
	md.diagnostic.referential=float('nan')*numpy.ones((md.mesh.numberofvertices,6))

	#Dirichlet Values
	if isinstance(md.inversion.vx_obs,numpy.ndarray) and numpy.size(md.inversion.vx_obs,axis=0)==md.mesh.numberofvertices and isinstance(md.inversion.vy_obs,numpy.ndarray) and numpy.size(md.inversion.vy_obs,axis=0)==md.mesh.numberofvertices:
		print "      boundary conditions for diagnostic model: spc set as observed velocities"
		md.diagnostic.spcvx[pos]=md.inversion.vx_obs[pos]
		md.diagnostic.spcvy[pos]=md.inversion.vy_obs[pos]
	else:
		print "      boundary conditions for diagnostic model: spc set as zero"

	md.hydrology.spcwatercolumn=numpy.zeros((md.mesh.numberofvertices,2))
	pos=numpy.nonzero(md.mesh.vertexonboundary)[0]
	md.hydrology.spcwatercolumn[pos,0]=1

	#segment on Neumann (Ice Front)
#	pos=find(vertexonicefront(md.mesh.segments(:,1)) | vertexonicefront(md.mesh.segments(:,2)));
	pos=numpy.nonzero(numpy.logical_or(vertexonicefront[md.mesh.segments[:,0].astype(int)-1],vertexonicefront[md.mesh.segments[:,1].astype(int)-1]))[0]
	if   md.mesh.dimension==2:
		pressureload=md.mesh.segments[pos,:]
	elif md.mesh.dimension==3:
#		pressureload_layer1=[md.mesh.segments(pos,1:2)  md.mesh.segments(pos,2)+md.mesh.numberofvertices2d  md.mesh.segments(pos,1)+md.mesh.numberofvertices2d  md.mesh.segments(pos,3)];
		pressureload_layer1=numpy.hstack((md.mesh.segments[pos,0:2],md.mesh.segments[pos,1]+md.mesh.numberofvertices2d,md.mesh.segments[pos,0]+md.mesh.numberofvertices2d,md.mesh.segments[pos,2]))
		pressureload=numpy.zeros((0,5))
		for i in xrange(1,md.mesh.numberoflayers):
#			pressureload=[pressureload ;pressureload_layer1(:,1:4)+(i-1)*md.mesh.numberofvertices2d pressureload_layer1(:,5)+(i-1)*md.mesh.numberofelements2d ];
			pressureload=numpy.vstack((pressureload,numpy.hstack((pressureload_layer1[:,0:4]+(i-1)*md.mesh.numberofvertices2d,pressureload_layer1[:,4]+(i-1)*md.mesh.numberofelements2d))))

	#Add water or air enum depending on the element
#	pressureload=[pressureload 1*md.mask.elementonfloatingice(pressureload(:,end))+ 0*md.mask.elementongroundedice(pressureload(:,end))];
	pressureload=numpy.hstack((pressureload,1.*md.mask.elementonfloatingice[pressureload[:,-1].astype('int')-1].reshape(-1,1)+0.*md.mask.elementongroundedice[pressureload[:,-1].astype('int')-1].reshape(-1,1)))

	#plug onto model
	md.diagnostic.icefront=pressureload

	#Create zeros basalforcings and surfaceforcings
	if numpy.all(numpy.isnan(md.surfaceforcings.precipitation)) and (md.surfaceforcings.ispdd==1):
		md.surfaceforcings.precipitation=numpy.zeros((md.mesh.numberofvertices,1))
		print "      no surfaceforcings.precipitation specified: values set as zero"
	if numpy.all(numpy.isnan(md.surfaceforcings.mass_balance)) and (md.surfaceforcings.ispdd==0):
		md.surfaceforcings.mass_balance=numpy.zeros((md.mesh.numberofvertices,1))
		print "      no surfaceforcings.mass_balance specified: values set as zero"
	if numpy.all(numpy.isnan(md.basalforcings.melting_rate)):
		md.basalforcings.melting_rate=numpy.zeros((md.mesh.numberofvertices,1))
		print "      no basalforcings.melting_rate specified: values set as zero"
	if numpy.all(numpy.isnan(md.balancethickness.thickening_rate)):
		md.balancethickness.thickening_rate=numpy.zeros((md.mesh.numberofvertices,1))
		print "      no balancethickness.thickening_rate specified: values set as zero"

	md.prognostic.spcthickness=float('nan')*numpy.ones((md.mesh.numberofvertices,1))
	md.balancethickness.spcthickness=float('nan')*numpy.ones((md.mesh.numberofvertices,1))

	if isinstance(md.initialization.temperature,numpy.ndarray) and numpy.size(md.initialization.temperature,axis=0)==md.mesh.numberofvertices:
		md.thermal.spctemperature=float('nan')*numpy.ones((md.mesh.numberofvertices,1))
#		pos=find(md.mesh.vertexonsurface); md.thermal.spctemperature(pos)=md.initialization.temperature(pos); %impose observed temperature on surface
		pos=numpy.nonzero(md.mesh.vertexonsurface)[0]
		md.thermal.spctemperature[pos]=md.initialization.temperature[pos]    #impose observed temperature on surface
		if not isinstance(md.basalforcings.geothermalflux,numpy.ndarray) or not numpy.size(md.basalforcings.geothermalflux,axis=0)==md.mesh.numberofvertices:
			md.basalforcings.geothermalflux=numpy.zeros((md.mesh.numberofvertices,1))
			md.basalforcings.geothermalflux[numpy.nonzero(md.mask.vertexongroundedice)]=50.*10.**-3    #50mW/m2
	else:
		print "      no thermal boundary conditions created: no observed temperature found"

	return md

