import numpy
from FlagElements import *

def setmask(md, floatingicename, groundedicename):
	"""
	SETMASK - establish boundaries between grounded and floating ice.

	   By default, ice is considered grounded. The contour floatingicename defines nodes 
	   for which ice is floating. The contour groundedicename defines nodes inside an floatingice, 
	   that are grounded (ie: ice rises, islands, etc ...)
	   All input files are in the Argus format (extension .exp).

	   Usage:
	      md=setmask(md,floatingicename,groundedicename)

	   Examples:
	      md=setmask(md,'all','');
	      md=setmask(md,'Iceshelves.exp','Islands.exp');
	"""

	#Get assigned fields
	x = md.mesh.x
	y = md.mesh.y
	elements = md.mesh.elements

	#Assign elementonfloatingice, elementongroundedice, vertexongroundedice and vertexonfloatingice. Only change at your own peril! This is synchronized heavily with the GroundingLineMigration module. {{{
	elementonfloatingice = FlagElements(md, floatingicename)
	elementongroundedice = FlagElements(md, groundedicename) 

	#Because groundedice nodes and elements can be included into an floatingice, we need to update. Remember, all the previous 
	#arrays come from domain outlines that can intersect one another: 

	elementonfloatingice = numpy.logical_and(elementonfloatingice,numpy.logical_not(elementongroundedice))
	elementongroundedice = numpy.logical_not(elementonfloatingice)

	#the order here is important. we choose vertexongroundedice as default on the grounding line.
	vertexonfloatingice = numpy.zeros(md.mesh.numberofvertices,'bool')
	vertexongroundedice = numpy.zeros(md.mesh.numberofvertices,'bool')
	vertexongroundedice[md.mesh.elements[numpy.nonzero(elementongroundedice),:].astype(int)-1]=True
	vertexonfloatingice[numpy.nonzero(numpy.logical_not(vertexongroundedice))]=True
	#}}}

	#Return: 
	md.mask.elementonfloatingice = elementonfloatingice.astype(float)
	md.mask.vertexonfloatingice = vertexonfloatingice.astype(float)
	md.mask.elementongroundedice = elementongroundedice.astype(float)
	md.mask.vertexongroundedice = vertexongroundedice.astype(float)
	md.mask.vertexonwater = numpy.zeros(md.mesh.numberofvertices)
	md.mask.elementonwater = numpy.zeros(md.mesh.numberofelements)

	return md
