#Test Name: EarthSlr Dakota Sampling glaciers
import numpy as np

from dmeth_params_set import *
from gmshplanet import *
from gmtmask import *
from lovenumbers import *
from materials import *
from model import *
from nodalvalue import *
from normal_uncertain import *
from response_function import *
from solve import *


# Mesh earth
md = model()
md.cluster = generic('name', oshostname(), 'np', 5)
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 = 'test2006'

# 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

# Uncertainty quantification
# Ice sheets #{{{
npart = 1
nt = 1
partition = -np.ones((md.mesh.numberofelements, 1))
pos = np.where(late < -80)[0]
partition[pos] = 0
pos = np.where(np.logical_and.reduce((late > 70, late < 80, longe > -60, longe < -30)))[0]
partition[pos] = 0

# Variables
qmuvar = OrderedStruct()
qmuvar.surfaceload = normal_uncertain.normal_uncertain(
    'descriptor',   'scaled_SurfaceloadIceThicknessChange',
    'mean',         1 * np.ones((npart, nt)),
    'stddev',       1 * np.ones((npart, nt)), # 10% standard deviation
    'partition',    partition,
    'transient',    'on',
    'nsteps',       nt
)
#}}}

# Correlation
md.qmu.correlation_matrix = []

# Variables final declaration
md.qmu.variables = OrderedStruct()
md.qmu.variables.surfaceload = qmuvar.surfaceload

locations = np.array([1, 5, 10, 15, 20])

# Responses #{{{
md.qmu.responses.sealevel1 = response_function.response_function('descriptor', 'Outputdefinition1')
md.qmu.responses.sealevel2 = response_function.response_function('descriptor', 'Outputdefinition2')
md.qmu.responses.sealevel3 = response_function.response_function('descriptor', 'Outputdefinition3')
md.qmu.responses.sealevel4 = response_function.response_function('descriptor', 'Outputdefinition4')
md.qmu.responses.sealevel5 = response_function.response_function('descriptor', 'Outputdefinition5')

# Output definitions
for i in range(len(locations)):
    if i == 0:
        md.outputdefinition.definitions = [
            nodalvalue(
                'name',             'SNode',
                'definitionstring', 'Outputdefinition1',
                'model_string',     'Sealevel',
                'node',             locations[i]
            )
        ]
    else:
        md.outputdefinition.definitions.append(
            nodalvalue(
                'name',             'SNode',
                'definitionstring', 'Outputdefinition' + str(i + 1),
                'model_string',     'Sealevel',
                'node',             locations[i]
            )
        )
#}}}

# Algorithm #{{{
md.qmu.method = dakota_method.dakota_method('nond_samp')
md.qmu.method = dmeth_params_set(
    md.qmu.method,
    'seed',             1234,
    'samples',          10,
    'sample_type',      'random'
)
md.qmu.output = 1
#}}}
# Parameters #{{{
md.qmu.params.direct = True
md.qmu.params.interval_type = 'forward'
md.qmu.params.analysis_driver = 'matlab'
md.qmu.params.evaluation_scheduling = 'master'
md.qmu.params.processors_per_evaluation = 2
md.qmu.params.tabular_graphics_data = True
md.qmu.isdakota = 1
md.verbose = verbose(0)
md.verbose.qmu = 1
#}}}
# QMU statistics #{{{

# TODO: Abstract reshaping of arrays away to src/m/classes/qmustatistics.py::marshall or src/m/solve/WriteData.py
#
md.qmu.statistics.nfiles_per_directory = 2
md.qmu.statistics.ndirectories = 5

md.qmu.statistics.method[0]['name'] = 'Histogram'
md.qmu.statistics.method[0]['fields'] = ['Sealevel', 'BslrIce']
md.qmu.statistics.method[0]['steps'] = np.arange(1, 10 + 1).reshape(1, -1)
md.qmu.statistics.method[0]['nbins'] = 20

md.qmu.statistics.addmethod()
md.qmu.statistics.method[1]['name'] = 'MeanVariance'
md.qmu.statistics.method[1]['fields'] = ['Sealevel', 'BslrIce']
md.qmu.statistics.method[1]['steps'] = np.arange(1, 10 + 1).reshape(1, -1)

md.qmu.statistics.addmethod()
md.qmu.statistics.method[2]['name'] = 'SampleSeries'
md.qmu.statistics.method[2]['fields'] = ['Sealevel', 'BslrIce']
md.qmu.statistics.method[2]['steps'] = np.arange(1, 10 + 1).reshape(1, -1)
md.qmu.statistics.method[2]['indices'] = locations.reshape(1, -1)
#}}}

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

# Run without statistics computations
md.qmu.statistics.method[0]['name'] = 'None'
md = solve(md, 'Transient')

# Compare statistics with our own here
svalues = mds.results.StatisticsSolution[-1].SealevelSamples # all values at locations

dvalues = np.zeros((md.qmu.method.params.samples, len(locations)))
for i in range(md.qmu.method.params.samples):
    dvalues[i, :] = md.results.dakota.modelresults[i].TransientSolution[-1].Sealevel[locations - 1].flatten()

samplesnorm = np.linalg.norm(dvalues - svalues, 'fro')

# Fields and tolerances to track changes
field_names = ['Samples Norm']
field_tolerances = [1e-13]
field_values = [samplesnorm]
