#Test Name: MISMIP3DHO
import numpy as np
from model import *
from socket import gethostname
from triangle import *
from setmask import *
from parameterize import *
from setflowequation import *
from solve import *

md = triangle(model(),'../Exp/Square.exp',100000.)
md = setmask(md,'../Exp/SquareShelf.exp','')
md = parameterize(md,'../Par/SquareSheetShelf.py')
md.initialization.vx[:] = 1.
md.initialization.vy[:] = 1.
md.geometry.thickness[:] = 500. - md.mesh.x / 10000.
md.geometry.bed = -100. - md.mesh.x / 1000.
md.geometry.base = -md.geometry.thickness * md.materials.rho_ice / md.materials.rho_water
md.mask.groundedice_levelset = md.geometry.thickness + md.materials.rho_water / md.materials.rho_ice * md.geometry.bed
pos = np.where(md.mask.groundedice_levelset >= 0)
md.geometry.base[pos] = md.geometry.bed[pos]
md.geometry.surface = md.geometry.base + md.geometry.thickness
md = md.extrude(4,1.)
md = setflowequation(md,'HO','all')

#Boundary conditions:
md.mask.ice_levelset = -np.ones((md.mesh.numberofvertices,))
md.mask.ice_levelset[np.where(md.mesh.x == max(md.mesh.x))] = 0.
md.stressbalance.spcvx[:] = float('Nan')
md.stressbalance.spcvy[:] = float('Nan')
md.stressbalance.spcvz[:] = float('Nan')
posA = np.intersect1d(np.array(np.where(md.mesh.y < 1000000.1)),np.array(np.where(md.mesh.y > 999999.9)))
posB = np.intersect1d(np.array(np.where(md.mesh.y < 0.1)),np.array(np.where(md.mesh.y > -0.1)))
pos = np.unique(np.concatenate((posA,posB)))
md.stressbalance.spcvy[pos] = 0.
pos2 = np.intersect1d(np.array(np.where(md.mesh.x < 0.1)), np.array(np.where(md.mesh.x > -0.1)))
md.stressbalance.spcvx[pos2] = 0.
md.stressbalance.spcvy[pos2] = 0.

md.materials.rheology_B = 1. / ((10**-25)**(1./3.)) * np.ones((md.mesh.numberofvertices,))
md.materials.rheology_law = 'None'
md.friction.coefficient[:] = np.sqrt(1e7) * np.ones((md.mesh.numberofvertices,))
md.friction.p = 3. * np.ones((md.mesh.numberofelements,))
md.smb.mass_balance[:] = 1.
md.basalforcings.groundedice_melting_rate[:] = 0.
md.basalforcings.floatingice_melting_rate[:] = 30.
md.transient.isthermal = 0
md.transient.isstressbalance = 1
md.transient.isgroundingline = 1
md.transient.ismasstransport = 1
md.transient.issmb = 1
md.transient.requested_outputs = ['default','BasalforcingsFloatingiceMeltingRate']
md.groundingline.migration = 'SubelementMigration'
md.timestepping.final_time = 30
md.timestepping.time_step = 10

md.cluster = generic('name',gethostname(),'np',3)
md = solve(md,'Transient')

#Fields and tolerances to track changes
field_names     = [
	'Bed1','Surface1','Thickness1','Floatingice1','Vx1','Vy1','Vz1','Pressure1','FloatingiceMeltingrate1',
	'Bed2','Surface2','Thickness2','Floatingice2','Vx2','Vy2','Vz2','Pressure2','FloatingiceMeltingrate2',
	'Bed3','Surface3','Thickness3','Floatingice3','Vx3','Vy3','Vz3','Pressure3','FloatingiceMeltingrate3']
field_tolerances = [
	2e-11,5e-12,2e-11,1e-11,7e-10,3e-08,6e-10,1e-13,1e-13,
	3e-11,3e-11,9e-10,7e-11,9e-09,2e-07,1e-09,1e-10,1e-13,
	1e-9,2e-08,7e-09,2e-7 ,1e-03,8e-04,2e-09,1e-10,1e-13]
field_values = [
	md.results.TransientSolution[0].Base,
	md.results.TransientSolution[0].Surface,
	md.results.TransientSolution[0].Thickness,
	md.results.TransientSolution[0].MaskGroundediceLevelset,
	md.results.TransientSolution[0].Vx,
	md.results.TransientSolution[0].Vy,
	md.results.TransientSolution[0].Vz,
	md.results.TransientSolution[0].Pressure,
	md.results.TransientSolution[0].BasalforcingsFloatingiceMeltingRate,
	md.results.TransientSolution[1].Base,
	md.results.TransientSolution[1].Surface,
	md.results.TransientSolution[1].Thickness,
	md.results.TransientSolution[1].MaskGroundediceLevelset,
	md.results.TransientSolution[1].Vx,
	md.results.TransientSolution[1].Vy,
	md.results.TransientSolution[1].Vz,
	md.results.TransientSolution[1].Pressure,
	md.results.TransientSolution[1].BasalforcingsFloatingiceMeltingRate,
	md.results.TransientSolution[2].Base,
	md.results.TransientSolution[2].Surface,
	md.results.TransientSolution[2].Thickness,
	md.results.TransientSolution[2].MaskGroundediceLevelset,
	md.results.TransientSolution[2].Vx,
	md.results.TransientSolution[2].Vy,
	md.results.TransientSolution[2].Vz,
	md.results.TransientSolution[2].Pressure,
	md.results.TransientSolution[2].BasalforcingsFloatingiceMeltingRate,
	]