#Test Name: MomentOfIntertia
import numpy as np
from model import *
from socket import gethostname
from solve import *
from gmshplanet import *
from love_numbers import *
from paterson import *
from maskpsl import *
from gmtmask import *

#mesh earth:
md = model()
rad_e = 6.371012 * 1e3  # mean radius of Earth, km
md.mesh = gmshplanet('radius', rad_e, 'resolution', 1000.0)  # km resolution

#parameterize slr solution:
#slr loading:  {{{
late = sum(md.mesh.lat[md.mesh.elements - 1], 1) / 3
longe = sum(md.mesh.long[md.mesh.elements - 1], 1) / 3

md.slr.deltathickness = np.zeros((md.mesh.numberofelements, ))
pos = np.intersect1d(np.array(np.where(late < -75)), np.array(np.where(longe > 0)))
#python does not include last element in array slices, (6:7) -> [5:7]
md.slr.deltathickness[pos[5:7]] = -1

md.slr.sealevel = np.zeros((md.mesh.numberofvertices, ))
md.dsl.global_average_thermosteric_sea_level_change=np.zeros((2, ))
md.dsl.sea_surface_height_change_above_geoid=np.zeros((md.mesh.numberofvertices+1, ))
md.dsl.sea_water_pressure_change_at_sea_floor=zeros((md.mesh.numberofvertices+1, ))

md.slr.ocean_area_scaling = 1

#elastic loading from love numbers:
nlov = 1000
md.slr.love_h = np.array(love_numbers('h'))
md.slr.love_h = np.resize(md.slr.love_h, nlov + 1)
md.slr.love_k = np.array(love_numbers('k'))
md.slr.love_k = np.resize(md.slr.love_k, nlov + 1)
md.slr.love_l = np.array(love_numbers('l'))
md.slr.love_l = np.resize(md.slr.love_l, nlov + 1)

#}}}
#mask:  {{{
md.mask = maskpsl()  # use maskpsl class (instead of mask) to store the ocean function as a ocean_levelset
mask = gmtmask(md.mesh.lat, md.mesh.long)

icemask = np.ones((md.mesh.numberofvertices, ))
pos = np.where(mask == 0)
icemask[pos[0]] = -1
pos = np.where(sum(mask[md.mesh.elements - 1], 1) < 3)
icemask[md.mesh.elements[pos, :] - 1] = -1
md.mask.ice_levelset = icemask
md.mask.ocean_levelset = np.zeros((md.mesh.numberofvertices, ))
pos = np.where(md.mask.ice_levelset == 1)
md.mask.ocean_levelset[pos] = 1

#make sure that the ice level set is all inclusive:
md.mask.land_levelset = np.zeros((md.mesh.numberofvertices, ))
md.mask.groundedice_levelset = -np.ones((md.mesh.numberofvertices, ))

#make sure that the elements that have loads are fully grounded:
pos = np.nonzero(md.slr.deltathickness)[0]
md.mask.groundedice_levelset[md.mesh.elements[pos, :] - 1] = 1

#make sure wherever there is an ice load, that the mask is set to ice:
icemask[md.mesh.elements[pos, :] - 1] = -1
md.mask.ice_levelset = icemask
# }}}
#geometry {{{
di = md.materials.rho_ice / md.materials.rho_water
md.geometry.thickness = np.ones((md.mesh.numberofvertices, ))
md.geometry.surface = (1 - di) * np.zeros((md.mesh.numberofvertices, ))
md.geometry.base = md.geometry.surface - md.geometry.thickness
md.geometry.bed = md.geometry.base
# }}}
#materials {{{
md.initialization.temperature = 273.25 * np.ones((md.mesh.numberofvertices, ))
md.materials.rheology_B = paterson(md.initialization.temperature)
md.materials.rheology_n = 3 * np.ones((md.mesh.numberofelements, ))
# }}}
#Miscellaneous {{{
md.miscellaneous.name = 'test2010'
# }}}
#Solution parameters {{{
md.slr.reltol = float('NaN')
md.slr.abstol = 1e-3
md.slr.geodetic = 1
# }}}

#New stuff
md.slr.spcthickness = np.nan * np.ones((md.mesh.numberofvertices, ))
md.slr.Ngia = np.zeros((md.mesh.numberofvertices, ))
md.slr.Ugia = np.zeros((md.mesh.numberofvertices, ))
md.slr.hydro_rate = np.zeros((md.mesh.numberofvertices, ))

#eustatic + rigid + elastic run:
md.slr.rigid = 1
md.slr.elastic = 1
md.slr.rotation = 1
md.cluster = generic('name', gethostname(), 'np', 3)

# uncomment following 2 lines for
md = solve(md, 'Sealevelrise')
eus = md.results.SealevelriseSolution.SealevelRSLEustatic
slr = md.results.SealevelriseSolution.Sealevel
moixz = md.results.SealevelriseSolution.SealevelInertiaTensorXZ
moiyz = md.results.SealevelriseSolution.SealevelInertiaTensorYZ
moizz = md.results.SealevelriseSolution.SealevelInertiaTensorZZ

# analytical moi = > just checking FOR ICE only!!! {{{
# ...have to mute**slr induced MOI in Tria.cpp**prior to the comparison
#rad_e = rad_e * 1e3  # now in meters
#areas = GetAreasSphericalTria(md.mesh.elements, md.mesh.lat, md.mesh.long, rad_e)
#lat = late * pi / 180 lon = longe * pi / 180
#moi_xz = sum(-md.materials.rho_freshwater. * md.slr.deltathickness. * areas. * rad_e^2. * sin(lat). * cos(lat). * cos(lon))
#moi_yz = sum(-md.materials.rho_freshwater. * md.slr.deltathickness. * areas. * rad_e^2. * sin(lat). * cos(lat). * sin(lon))
# }}}

#Fields and tolerances to track changes
field_names = ['eus', 'slr', 'moixz', 'moiyz', 'moizz']
field_tolerances = [1e-13, 1e-13, 1e-13, 1e-13, 1e-13]
field_values = [eus, slr, moixz, moiyz, moizz]
