from fielddisplay import fielddisplay
from MatlabFuncs import *
from model import *
import numpy as np
from checkfield import checkfield
from WriteData import WriteData


class slr(object):
    """
    SLR class definition

        Usage:
          slr = slr()
    """
    def __init__(self):  # {{{
        self.deltathickness = float('NaN')
        self.sealevel = float('NaN')
        self.spcthickness = float('NaN')
        self.maxiter = 0
        self.reltol = 0
        self.abstol = 0
        self.love_h = 0  #provided by PREM model()
        self.love_k = 0  #ideam
        self.love_l = 0  #ideam
        self.tide_love_k = 0  #ideam
        self.tide_love_h = 0  #ideam
        self.fluid_love = 0
        self.equatorial_moi = 0
        self.polar_moi = 0
        self.angular_velocity = 0
        self.rigid = 0
        self.elastic = 0
        self.rotation = 0
        self.ocean_area_scaling = 0
        self.steric_rate = 0  #rate of ocean expansion from steric effects.
        self.hydro_rate = 0  #rate of steric expansion from hydrological effects.
        self.geodetic_run_frequency = 1  #how many time steps we skip before we run the geodetic part of the solver during transient
        self.geodetic = 0  #compute geodetic SLR? (in addition to steric?)
        self.degacc = 0
        self.loop_increment = 0
        self.horiz = 0
        self.Ngia = float('NaN')
        self.Ugia = float('NaN')
        self.requested_outputs = []
        self.transitions = []

        #set defaults
        self.setdefaultparameters()
    #}}}

    def __repr__(self):  # {{{
        string = '   slr parameters:'
        string = "%s\n%s" % (string, fielddisplay(self, 'deltathickness', 'thickness change: ice height equivalent [m]'))
        string = "%s\n%s" % (string, fielddisplay(self, 'sealevel', 'current sea level (prior to computation) [m]'))
        string = "%s\n%s" % (string, fielddisplay(self, 'spcthickness', 'thickness constraints (NaN means no constraint) [m]'))
        string = "%s\n%s" % (string, fielddisplay(self, 'reltol', 'sea level rise relative convergence criterion, (NaN: not applied)'))
        string = "%s\n%s" % (string, fielddisplay(self, 'abstol', 'sea level rise absolute convergence criterion, (default, NaN: not applied)'))
        string = "%s\n%s" % (string, fielddisplay(self, 'maxiter', 'maximum number of nonlinear iterations'))
        string = "%s\n%s" % (string, fielddisplay(self, 'love_h', 'load Love number for radial displacement'))
        string = "%s\n%s" % (string, fielddisplay(self, 'love_k', 'load Love number for gravitational potential perturbation'))
        string = "%s\n%s" % (string, fielddisplay(self, 'love_l', 'load Love number for horizontal displaements'))
        string = "%s\n%s" % (string, fielddisplay(self, 'tide_love_k', 'tidal load Love number (degree 2)'))
        string = "%s\n%s" % (string, fielddisplay(self, 'tide_love_h', 'tidal load Love number (degree 2)'))
        string = "%s\n%s" % (string, fielddisplay(self, 'fluid_love', 'secular fluid Love number'))
        string = "%s\n%s" % (string, fielddisplay(self, 'equatorial_moi', 'mean equatorial moment of inertia [kg m^2]'))
        string = "%s\n%s" % (string, fielddisplay(self, 'polar_moi', 'polar moment of inertia [kg m^2]'))
        string = "%s\n%s" % (string, fielddisplay(self, 'angular_velocity', 'mean rotational velocity of earth [per second]'))
        string = "%s\n%s" % (string, fielddisplay(self, 'ocean_area_scaling', 'correction for model representation of ocean area [default: No correction]'))
        string = "%s\n%s" % (string, fielddisplay(self, 'steric_rate', 'rate of steric ocean expansion [mm / yr]'))
        string = "%s\n%s" % (string, fielddisplay(self, 'hydro_rate', 'rate of hydrological expansion [mm / yr]'))
        string = "%s\n%s" % (string, fielddisplay(self, 'Ngia', 'rate of viscous (GIA) geoid expansion (in mm / yr)'))
        string = "%s\n%s" % (string, fielddisplay(self, 'Ugia', 'rate of viscous (GIA) bedrock uplift (in mm / yr)'))
        string = "%s\n%s" % (string, fielddisplay(self, 'loop_increment', 'vector assembly (in the convolution) framentation'))
        string = "%s\n%s" % (string, fielddisplay(self, 'geodetic', 'compute geodetic SLR? (in addition to steric?) default 0'))
        string = "%s\n%s" % (string, fielddisplay(self, 'geodetic_run_frequency', 'how many time steps we skip before we run SLR solver during transient (default: 1)'))
        string = "%s\n%s" % (string, fielddisplay(self, 'rigid', 'rigid earth graviational potential perturbation'))
        string = "%s\n%s" % (string, fielddisplay(self, 'elastic', 'elastic earth graviational potential perturbation'))
        string = "%s\n%s" % (string, fielddisplay(self, 'rotation', 'earth rotational potential perturbation'))
        string = "%s\n%s" % (string, fielddisplay(self, 'degacc', 'accuracy (default .01 deg) for numerical discretization of the Green''s functions'))
        string = "%s\n%s" % (string, fielddisplay(self, 'transitions', 'indices into parts of the mesh that will be icecaps'))
        string = "%s\n%s" % (string, fielddisplay(self, 'requested_outputs', 'additional outputs requested'))

        return string
    # }}}

    def setdefaultparameters(self):  # {{{
        #Convergence criterion: absolute, relative and residual
        self.reltol = 0.01  #default
        self.abstol = float('NaN')  #1 mm of sea level rise
        #maximum of non - linear iterations.
        self.maxiter = 5
        self.loop_increment = 200
        #computational flags:
        self.geodetic = 0
        self.rigid = 1
        self.elastic = 1
        self.ocean_area_scaling = 0
        self.rotation = 1
        #tidal love numbers:
        self.tide_love_h = 0.6149  #degree 2
        self.tide_love_k = 0.3055  #degree 2
        #secular fluid love number:
        self.fluid_love = 0.942
        #moment of inertia:
        self.equatorial_moi = 8.0077 * 1.0e37  # [kg m^2]
        self.polar_moi = 8.0345 * 1.0e37  # [kg m^2]
        #mean rotational velocity of earth
        self.angular_velocity = 7.2921 * 1.0e-5  # [s^ - 1]
        #numerical discretization accuracy
        self.degacc = 0.01
        #steric:
        self.steric_rate = 0
        self.hydro_rate = 0
        #how many time steps we skip before we run SLR solver during transient
        self.geodetic_run_frequency = 1
        #output default:
        self.requested_outputs = ['default']
        #transitions should be a cell array of vectors:
        self.transitions = []
        #horizontal displacement?  (not by default)
        self.horiz = 0

        return self
    #}}}

    def checkconsistency(self, md, solution, analyses):  # {{{
        #Early return
        if (solution != 'SealevelriseAnalysis') or (solution == 'TransientSolution' and not md.transient.isslr):
            return md

        md = checkfield(md, 'fieldname', 'slr.deltathickness', 'NaN', 1, 'Inf', 1, 'size', [md.mesh.numberofelements])
        md = checkfield(md, 'fieldname', 'slr.sealevel', 'NaN', 1, 'Inf', 1, 'size', [md.mesh.numberofvertices])
        md = checkfield(md, 'fieldname', 'slr.spcthickness', 'Inf', 1, 'timeseries', 1)
        md = checkfield(md, 'fieldname', 'slr.love_h', 'NaN', 1, 'Inf', 1)
        md = checkfield(md, 'fieldname', 'slr.love_k', 'NaN', 1, 'Inf', 1)
        md = checkfield(md, 'fieldname', 'slr.love_l', 'NaN', 1, 'Inf', 1)
        md = checkfield(md, 'fieldname', 'slr.tide_love_h', 'NaN', 1, 'Inf', 1)
        md = checkfield(md, 'fieldname', 'slr.tide_love_k', 'NaN', 1, 'Inf', 1)
        md = checkfield(md, 'fieldname', 'slr.fluid_love', 'NaN', 1, 'Inf', 1)
        md = checkfield(md, 'fieldname', 'slr.equatorial_moi', 'NaN', 1, 'Inf', 1)
        md = checkfield(md, 'fieldname', 'slr.polar_moi', 'NaN', 1, 'Inf', 1)
        md = checkfield(md, 'fieldname', 'slr.angular_velocity', 'NaN', 1, 'Inf', 1)
        md = checkfield(md, 'fieldname', 'slr.reltol', 'size', [1, 1])
        md = checkfield(md, 'fieldname', 'slr.abstol', 'size', [1, 1])
        md = checkfield(md, 'fieldname', 'slr.maxiter', 'size', [1, 1], '>=', 1)
        md = checkfield(md, 'fieldname', 'slr.geodetic_run_frequency', 'size', [1, 1], '>=', 1)
        md = checkfield(md, 'fieldname', 'slr.steric_rate', 'NaN', 1, 'Inf', 1, 'size', [md.mesh.numberofvertices])
        md = checkfield(md, 'fieldname', 'slr.hydro_rate', 'NaN', 1, 'Inf', 1, 'size', [md.mesh.numberofvertices])
        md = checkfield(md, 'fieldname', 'slr.degacc', 'size', [1, 1], '>=', 1e-10)
        md = checkfield(md, 'fieldname', 'slr.requested_outputs', 'stringrow', 1)
        md = checkfield(md, 'fieldname', 'slr.loop_increment', 'NaN', 1, 'Inf', 1, '>=', 1)
        md = checkfield(md, 'fieldname', 'slr.horiz', 'NaN', 1, 'Inf', 1, 'values', [0, 1])
        md = checkfield(md, 'fieldname', 'slr.Ngia', 'NaN', 1, 'Inf', 1, 'size', [md.mesh.numberofvertices])
        md = checkfield(md, 'fieldname', 'slr.Ugia', 'NaN', 1, 'Inf', 1, 'size', [md.mesh.numberofvertices])

        #check that love numbers are provided at the same level of accuracy:
        if (size(self.love_h, 0) != size(self.love_k, 0) | size(self.love_h, 0) != size(self.love_l, 0)):
            error('slr error message: love numbers should be provided at the same level of accuracy')

        #cross check that whereever we have an ice load, the mask is < 0 on each vertex:
        pos = np.where(self.deltathickness)
        maskpos = md.mask.ice_levelset[md.mesh.elements[pos, :]]
        els = np.where(maskpos > 0)
        if len(els[0]) > 0:
            warnings.warn('slr checkconsistency fail: there are elements with ice loads where some vertices are not on the ice!')

        #check that  if geodetic is requested, we are a mesh3dsurface model (planet), or if we are not,
        #a coupler to a planet model is provided.
        if self.geodetic and not md.transient.iscoupler and domaintype(md.mesh) != 'mesh3dsurface':
            error('model is requesting geodetic computations without being a mesh3dsurface, or being coupled to one!')
        return md
    # }}}

    def defaultoutputs(self, md):  # {{{
        return ['Sealevel']
    # }}}

    def marshall(self, prefix, md, fid):  # {{{
        WriteData(fid, prefix, 'object', self, 'fieldname', 'deltathickness', 'format', 'DoubleMat', 'mattype', 2)
        WriteData(fid, prefix, 'object', self, 'fieldname', 'sealevel', 'mattype', 1, 'format', 'DoubleMat', 'timeserieslength', md.mesh.numberofvertices + 1, 'yts', md.constants.yts)
        WriteData(fid, prefix, 'object', self, 'fieldname', 'spcthickness', 'format', 'DoubleMat', 'mattype', 1, 'timeserieslength', md.mesh.numberofvertices + 1, 'yts', md.constants.yts)
        WriteData(fid, prefix, 'object', self, 'fieldname', 'reltol', 'format', 'Double')
        WriteData(fid, prefix, 'object', self, 'fieldname', 'abstol', 'format', 'Double')
        WriteData(fid, prefix, 'object', self, 'fieldname', 'maxiter', 'format', 'Integer')
        WriteData(fid, prefix, 'object', self, 'fieldname', 'love_h', 'format', 'DoubleMat', 'mattype', 1)
        WriteData(fid, prefix, 'object', self, 'fieldname', 'love_k', 'format', 'DoubleMat', 'mattype', 1)
        WriteData(fid, prefix, 'object', self, 'fieldname', 'love_l', 'format', 'DoubleMat', 'mattype', 1)
        WriteData(fid, prefix, 'object', self, 'fieldname', 'tide_love_h', 'format', 'Double')
        WriteData(fid, prefix, 'object', self, 'fieldname', 'tide_love_k', 'format', 'Double')
        WriteData(fid, prefix, 'object', self, 'fieldname', 'fluid_love', 'format', 'Double')
        WriteData(fid, prefix, 'object', self, 'fieldname', 'equatorial_moi', 'format', 'Double')
        WriteData(fid, prefix, 'object', self, 'fieldname', 'polar_moi', 'format', 'Double')
        WriteData(fid, prefix, 'object', self, 'fieldname', 'angular_velocity', 'format', 'Double')
        WriteData(fid, prefix, 'object', self, 'fieldname', 'rigid', 'format', 'Boolean')
        WriteData(fid, prefix, 'object', self, 'fieldname', 'elastic', 'format', 'Boolean')
        WriteData(fid, prefix, 'object', self, 'fieldname', 'rotation', 'format', 'Boolean')
        WriteData(fid, prefix, 'object', self, 'fieldname', 'ocean_area_scaling', 'format', 'Boolean')
        WriteData(fid, prefix, 'object', self, 'fieldname', 'geodetic_run_frequency', 'format', 'Integer')
        WriteData(fid, prefix, 'object', self, 'fieldname', 'steric_rate', 'format', 'DoubleMat', 'mattype', 1, 'scale', 1e-3 / md.constants.yts)
        WriteData(fid, prefix, 'object', self, 'fieldname', 'hydro_rate', 'format', 'DoubleMat', 'mattype', 1, 'scale', 1e-3 / md.constants.yts)
        WriteData(fid, prefix, 'object', self, 'fieldname', 'Ngia', 'format', 'DoubleMat', 'mattype', 1, 'scale', 1e-3 / md.constants.yts)
        WriteData(fid, prefix, 'object', self, 'fieldname', 'Ugia', 'format', 'DoubleMat', 'mattype', 1, 'scale', 1e-3 / md.constants.yts)
        WriteData(fid, prefix, 'object', self, 'fieldname', 'degacc', 'format', 'Double')
        WriteData(fid, prefix, 'object', self, 'fieldname', 'transitions', 'format', 'MatArray')
        WriteData(fid, prefix, 'object', self, 'fieldname', 'loop_increment', 'format', 'Integer')
        WriteData(fid, prefix, 'object', self, 'fieldname', 'horiz', 'format', 'Integer')
        WriteData(fid, prefix, 'object', self, 'fieldname', 'geodetic', 'format', 'Integer')

    #process requested outputs
        outputs = self.requested_outputs
        indices = [i for i, x in enumerate(outputs) if x == 'default']
        if len(indices) > 0:
            outputscopy = outputs[0:max(0, indices[0] - 1)] + self.defaultoutputs(md) + outputs[indices[0] + 1:]
            outputs = outputscopy
        WriteData(fid, prefix, 'data', outputs, 'name', 'md.slr.requested_outputs', 'format', 'StringArray')
    # }}}
