#Test Name: EarthSlr
import numpy as np

from gmshplanet import *
from gmtmask import *
from lovenumbers import *
from materials import *
from model import *
from parameterize import *
from paterson import *
from solve import *


# Mesh earth
md = model()
md.mesh = gmshplanet('radius', 6.371012 * 1e3, 'resolution', 700.) #700 km resolution mesh

# Parameterize solidearth solution
# Solidearth loading #{{{
md.solidearth.surfaceload.icethicknesschange = np.zeros((md.mesh.numberofelements, 1))
md.solidearth.initialsealevel = np.zeros((md.mesh.numberofvertices, 1))
md.dsl.global_average_thermosteric_sea_level_change = np.zeros((2, 1))
md.dsl.sea_surface_height_change_above_geoid = np.zeros((md.mesh.numberofvertices + 1, 1))
md.dsl.sea_water_pressure_change_at_sea_floor = np.zeros((md.mesh.numberofvertices + 1, 1))

# Antarctica
late = md.mesh.lat[md.mesh.elements - 1].sum(axis=1) / 3
longe = md.mesh.long[md.mesh.elements - 1].sum(axis=1) / 3
pos = np.where(late < -80)[0]
md.solidearth.surfaceload.icethicknesschange[pos] = -100
# Greenland
pos = np.where(np.logical_and.reduce((late > 70, late < 80, longe > -60, longe < -30)))[0]
md.solidearth.surfaceload.icethicknesschange[pos] = -100

# Elastic loading from love numbers
md.solidearth.lovenumbers = lovenumbers('maxdeg', 100)
#}}}

# Mask #{{{
mask = gmtmask(md.mesh.lat, md.mesh.long)
icemask = np.ones((md.mesh.numberofvertices, 1))
pos = np.where(mask == 0)[0]
icemask[pos] = -1
pos = np.where(mask[md.mesh.elements - 1].sum(axis=1) < 3)[0]
icemask[md.mesh.elements[pos, :] - 1] = -1
md.mask.ice_levelset = icemask
md.mask.ocean_levelset = -icemask

# Make sure that the elements that have loads are fully grounded
pos = np.nonzero(md.solidearth.surfaceload.icethicknesschange)[0]
md.mask.ocean_levelset[md.mesh.elements[pos, :] - 1] = 1

# Make sure wherever there is an ice load, that the mask is set to ice:
#pos = np.nonzero(md.solidearth.surfaceload.icethicknesschange)[0] # TODO: Do we need to do this twice?
md.mask.ice_levelset[md.mesh.elements[pos, :] - 1] = -1
# }}}

md.solidearth.settings.ocean_area_scaling = 0

# Geometry for the bed; arbitrary
md.geometry.bed = -np.ones((md.mesh.numberofvertices, 1))

# Materials
md.materials = materials('hydro')

# Miscellaneous
md.miscellaneous.name = 'test2005'

# Solution parameters
md.solidearth.settings.reltol = np.nan
md.solidearth.settings.abstol = 1e-3
md.solidearth.settings.computesealevelchange = 1

# Max number of iterations reverted back to 10 (i.e. the original default value)
md.solidearth.settings.maxiter = 10

# Eustatic + rigid + elastic + rotation run
md.solidearth.settings.rigid = 1
md.solidearth.settings.elastic = 1
md.solidearth.settings.rotation = 1

# Transient settings
md.timestepping.start_time = 0
md.timestepping.final_time = 10
md.timestepping.time_step = 1
md.transient.isslr = 1
md.transient.issmb = 0
md.transient.isgia = 1
md.transient.ismasstransport = 0
md.transient.isstressbalance = 0
md.transient.isthermal = 0
dh = np.asarray(md.solidearth.surfaceload.icethicknesschange).T
deltathickness = np.zeros((md.mesh.numberofelements + 1, 10))
for i in range(10):
    deltathickness[0:-1, i] = dh * (i + 1)
deltathickness[-1, :] = np.arange(0, 10, 1)
md.solidearth.surfaceload.icethicknesschange = deltathickness

# Hack
md.geometry.surface = np.zeros((md.mesh.numberofvertices, 1))
md.geometry.thickness = np.ones((md.mesh.numberofvertices, 1))
md.geometry.base = -np.ones((md.mesh.numberofvertices, 1))
md.geometry.bed = md.geometry.base

# Run transient solution
md = solve(md, 'Transient')

S1 = md.results.TransientSolution[1 - 1].Sealevel
S5 = md.results.TransientSolution[5 - 1].Sealevel
S10 = md.results.TransientSolution[10 - 1].Sealevel
Seus10 = md.results.TransientSolution[10 - 1].Bslr

#Fields and tolerances to track changes
field_names = ['Sealevel1', 'Sealevel5', 'Sealevel10', 'Seustatic10']
field_tolerances = [1e-13, 1e-13, 1e-13, 1e-13]
field_values = [S1, S5, S10, Seus10]
