#Test Name: SquareShelfSMBautoregression
from __future__ import division  # TODO: Remove this import after ISSM is updated to use Python 3 by default
import numpy as np
from model import *
from socket import gethostname
from parameterize import parameterize
from setflowequation import setflowequation
from setmask import setmask
from SMBautoregression import SMBautoregression
from solve import solve
from triangle import triangle


md = triangle(model(), '../Exp/Square.exp', 80000)
md = setmask(md, 'all', '')
md = parameterize(md, '../Par/SquareShelfConstrained.py')
md = setflowequation(md, 'SSA', 'all')
md.cluster = generic('name', gethostname(), 'np', 3)
md.transient.requested_outputs = ['default', 'IceVolume', 'SmbMassBalance']

ymax = max(md.mesh.y)
xmax = max(md.mesh.x)
# Generate basin IDs for 3 basins
idbasin = np.zeros((md.mesh.numberofelements,))
iid1 = np.where(md.mesh.y >= 2. / 3. * ymax)[0]
iid2 = np.where(np.logical_and(md.mesh.y < 2. / 3. * ymax, md.mesh.x >= 1. / 3. * xmax))[0]
iid3 = np.where(np.logical_and(md.mesh.y < 2. / 3. * ymax, md.mesh.x < 1. / 3. * xmax))[0]
for ii in range(md.mesh.numberofelements):
    for vertex in range(3):
        if md.mesh.elements[ii][vertex] - 1 in iid1:  # one vertex in basin 1; NOTE: offset because of 1-based vertex indexing
            idbasin[ii] = 1
    if idbasin[ii] == 0:  # no vertex was found in basin 1
        for vertex in range(3):
            if md.mesh.elements[ii][vertex] - 1 in iid2:  # one vertex in basin 2; NOTE: offset because of 1-based vertex indexing
                idbasin[ii] = 2
    if idbasin[ii] == 0:  # no vertex was found in basin 1 and 2
        idbasin[ii] = 3

# SMB parameters
md.timestepping.start_time = 0
md.timestepping.time_step = 1
md.timestepping.final_time = 5
md.smb = SMBautoregression()
md.smb.num_basins = 3  # number of basins
md.smb.basin_id = idbasin  # prescribe basin ID number to elements;
md.smb.beta0 = np.array([[0.5, 1.2, 1.5]])  # intercept values of SMB in basins [m ice eq./yr]
md.smb.beta1 = np.array([[0.0, 0.01, -0.01]])  # trend values of SMB in basins [m ice eq./yr^2]
md.smb.ar_initialtime = md.timestepping.start_time
md.smb.ar_order = 4
md.smb.ar_timestep = 2.0  #timestep of the autoregressive model [yr]
md.smb.phi = np.array([[0.2, 0.1, 0.05, 0.01], [0.4, 0.2, -0.2, 0.1], [0.4, -0.4, 0.1, -0.1]])
md.smb.lapserates        = np.array([[0.01,0.0],[0.01,-0.01],[0.0,-0.01]])
md.smb.elevationbins  = np.array([100,150,100]).reshape(md.smb.num_basins,1)

# Stochastic forcing
md.stochasticforcing.isstochasticforcing = 1
md.stochasticforcing.fields = ['SMBautoregression']
md.stochasticforcing.covariance = np.array([[0.15, 0.08, -0.02], [0.08, 0.12, -0.05], [-0.02, -0.05, 0.1]])  # global covariance among- and between-fields
md.stochasticforcing.randomflag = 0  # fixed random seeds
md.stochasticforcing.stochastictimestep  = 1.0

md = solve(md, 'Transient')

# Fields and tolerances to track changes
field_names = [
    'Vx1', 'Vy1', 'Vel1', 'Thickness1', 'Volume1', 'SmbMassBalance1',
    'Vx2', 'Vy2', 'Vel2', 'Thickness2', 'Volume2', 'SmbMassBalance2',
    'Vx3', 'Vy3', 'Vel3', 'Thickness3', 'Volume3', 'SmbMassBalance3'
]
field_tolerances = [
    1e-13, 1e-13, 1e-13, 1e-13, 1e-13, 1e-13,
    1e-13, 1e-13, 1e-13, 1e-13, 1e-13, 1e-13,
    1e-13, 1e-13, 1e-13, 1e-13, 1e-13, 1e-13
]
field_values = [
    md.results.TransientSolution[0].Vx,
    md.results.TransientSolution[0].Vy,
    md.results.TransientSolution[0].Vel,
    md.results.TransientSolution[0].Thickness,
    md.results.TransientSolution[0].IceVolume,
    md.results.TransientSolution[0].SmbMassBalance,
    md.results.TransientSolution[1].Vx,
    md.results.TransientSolution[1].Vy,
    md.results.TransientSolution[1].Vel,
    md.results.TransientSolution[1].Thickness,
    md.results.TransientSolution[1].IceVolume,
    md.results.TransientSolution[1].SmbMassBalance,
    md.results.TransientSolution[2].Vx,
    md.results.TransientSolution[2].Vy,
    md.results.TransientSolution[2].Vel,
    md.results.TransientSolution[2].Thickness,
    md.results.TransientSolution[2].IceVolume,
    md.results.TransientSolution[2].SmbMassBalance
]
