import numpy
from SetMarineIceSheetBC import *

print "      creating thickness"
hmin=0.01
hmax=2756.7
radius=numpy.sqrt((md.mesh.x)**2+(md.mesh.y)**2).reshape(-1,1)
radiusmax=numpy.max(radius)
radius[numpy.nonzero(radius>(1.-10**-9)*radiusmax)]=radiusmax    #eliminate roundoff issues in next statement
md.geometry.thickness=hmin*numpy.ones((numpy.size(md.mesh.x),1))+hmax*(4.*((1./2.)**(4./3.)*numpy.ones((numpy.size(md.mesh.x),1))-((radius)/(2.*radiusmax))**(4./3.)))**(3./8.)
md.geometry.bed=0.*md.geometry.thickness
md.geometry.surface=md.geometry.bed+md.geometry.thickness

print "      creating drag"
md.friction.coefficient=20.*numpy.ones((md.mesh.numberofvertices,1))    #q=1. no drag is specified in the analytical solution
#Take care of iceshelves: no basal drag
pos=numpy.nonzero(md.mask.elementonfloatingice)[0]
md.friction.coefficient[md.mesh.elements[pos,:]-1]=0.
md.friction.p=numpy.ones((md.mesh.numberofelements,1))
md.friction.q=numpy.ones((md.mesh.numberofelements,1))

print "      creating temperatures"
tmin=238.15    #K
st=1.67*10**-2/1000.    #k/m
md.initialization.temperature=tmin+st*radius
md.basalforcings.geothermalflux=4.2*10**-2*numpy.ones((md.mesh.numberofvertices,1))

print "      creating flow law parameter"
md.materials.rheology_B=6.81*10**7*numpy.ones((md.mesh.numberofvertices,1))    #to have the same B as the analytical solution 
md.materials.rheology_n=3.*numpy.ones((md.mesh.numberofelements,1))

print "      creating surface mass balance"
smb_max=0.5    #m/yr
sb=10**-2/1000.    #m/yr/m
rel=450.*1000.    #m
md.surfaceforcings.mass_balance=numpy.minimum(smb_max*numpy.ones_like(radius),sb*(rel-radius))

print "      creating velocities"
constant=0.3
md.inversion.vx_obs=constant/2.*md.mesh.x.reshape(-1,1)*(md.geometry.thickness)**-1
md.inversion.vy_obs=constant/2.*md.mesh.y.reshape(-1,1)*(md.geometry.thickness)**-1
md.inversion.vel_obs=numpy.sqrt((md.inversion.vx_obs)**2+(md.inversion.vy_obs)**2)
md.initialization.vx=numpy.zeros((md.mesh.numberofvertices,1))
md.initialization.vy=numpy.zeros((md.mesh.numberofvertices,1))
md.initialization.vz=numpy.zeros((md.mesh.numberofvertices,1))
md.initialization.pressure=numpy.zeros((md.mesh.numberofvertices,1))

#Deal with boundary conditions:
print "      boundary conditions for diagnostic model:"
md=SetMarineIceSheetBC(md,'../Exp/RoundFrontEISMINT.exp')

radius=numpy.sqrt((md.mesh.x)**2+(md.mesh.y)**2)
pos=numpy.nonzero(radius==numpy.min(radius))[0]
md.mesh.x[pos]=0.
md.mesh.y[pos]=0.    #the closest node to the center is changed to be exactly at the center

md.diagnostic.spcvx[pos]=0.
md.diagnostic.spcvy[pos]=0.
md.diagnostic.spcvz[pos]=0.
