import os
import numpy

def SetIceShelfBC(md,icefrontfile=''):
	"""
	SETICESHELFBC - Create the boundary conditions for diagnostic and thermal models for a  Ice Shelf with Ice Front

	   Neumann BC are used on the ice front (an ANRGUS contour around the ice front
	   must be given in input)
	   Dirichlet BC are used elsewhere for diagnostic

	   Usage:
	      md=SetIceShelfBC(md,varargin)

	   Example:
	      md=SetIceShelfBC(md);
	      md=SetIceShelfBC(md,'Front.exp');

	   See also: SETICESHEETBC, SETMARINEICESHEETBC
	"""

	#node on Dirichlet (boundary and ~icefront)
	if icefrontfile:
		if not os.path.exists(icefrontfile):
			raise IOError("SetIceShelfBC error message: ice front file '%s' not found." % icefrontfile)
		nodeinsideicefront=ContourToMesh(md.mesh.elements,md.mesh.x,md.mesh.y,icefrontfile,'node',2)
		nodeonicefront=double(md.mesh.vertexonboundary and nodeinsideicefront)
	else:
		nodeonicefront=numpy.zeros(md.mesh.numberofvertices)

#	pos=find(md.mesh.vertexonboundary & ~nodeonicefront);
	pos=[i for i,(vob,noif) in enumerate(zip(md.mesh.vertexonboundary,nodeonicefront)) if vob and not noif]
	md.diagnostic.spcvx=float('NaN')*numpy.ones(md.mesh.numberofvertices)
	md.diagnostic.spcvy=float('NaN')*numpy.ones(md.mesh.numberofvertices)
	md.diagnostic.spcvz=float('NaN')*numpy.ones(md.mesh.numberofvertices)
	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 numpy.size(md.inversion.vx_obs)==md.mesh.numberofvertices and numpy.size(md.inversion.vy_obs)==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'

	#segment on Ice Front
	#segment on Neumann (Ice Front)
#	pos=find(nodeonicefront(md.mesh.segments(:,1)) | nodeonicefront(md.mesh.segments(:,2)));
	pos=[i for i,(noif1,noif2) in enumerate(zip(nodeonicefront[md.mesh.segments[:,0].astype('int')-1],nodeonicefront[md.mesh.segments[:,1].astype('int')-1])) if noif1 or noif2]
	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.concatenate((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]),axis=1)
		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.concatenate((pressureload,numpy.concatenate((pressureload_layer1[:,0:3]+(i-1)*md.mesh.numberofvertices2d,pressureload_layer1[:,4]+(i-1)*md.mesh.numberofelements2d),axis=1)),axis=0)

	#Add water or air enum depending on the element
#	pressureload=[pressureload 1*md.mask.elementonfloatingice(pressureload(:,end))];
	pressureload=numpy.concatenate((pressureload,1*md.mask.elementonfloatingice[pressureload[:,-1].astype('int')-1].reshape((-1,1))),axis=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)
		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)
		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)
		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)
		print '      no balancethickness.thickening_rate specified: values set as zero'

	md.prognostic.spcthickness=float('NaN')*numpy.ones(md.mesh.numberofvertices)
	md.balancethickness.spcthickness=float('NaN')*numpy.ones(md.mesh.numberofvertices)

	if numpy.size(md.initialization.temperature)==md.mesh.numberofvertices:
		md.thermal.spctemperature=float('NaN')*numpy.ones(md.mesh.numberofvertices)
#		pos=find(md.mesh.vertexonsurface); md.thermal.spctemperature(pos)=md.initialization.temperature(pos); %impose observed temperature on surface
		pos=[i for i,vos in enumerate(md.mesh.vertexonsurface) if vos]
		md.thermal.spctemperature[pos]=md.initialization.temperature[pos]    # impose observed temperature on surface
		if not numpy.size(md.basalforcings.geothermalflux)==md.mesh.numberofvertices:
			md.basalforcings.geothermalflux=numpy.zeros(md.mesh.numberofvertices)
	else:
		print '      no thermal boundary conditions created: no observed temperature found'

	return md