# -*- coding: utf-8 -*-
import numpy as np
from checkfield import checkfield
from fielddisplay import fielddisplay
from MatlabFuncs import *
from WriteData import WriteData


class frontalforcingsrignotautoregression(object):
    """FRONTALFORCINGSRIGNOTAUTOREGRESSION class definition

    Usage:
        frontalforcingsrignotautoregression = frontalforcingsrignotautoregression()
    """

    def __init__(self, *args):  # {{{
        self.num_basins = 0
        self.beta0 = np.nan
        self.beta1 = np.nan
        self.ar_order = 0
        self.ar_initialtime = 0
        self.ar_timestep = 0
        self.phi = np.nan
        self.basin_id = np.nan
        self.subglacial_discharge = np.nan

        if len(args) == 0:
            self.setdefaultparameters()
        else:
            error('constructor not supported')

    def __repr__(self):  # {{{
        s = '   Frontalforcings parameters:\n'
        s += '{}\n'.format(fielddisplay(self, 'num_basins', 'number of different basins [unitless]'))
        s += '{}\n'.format(fielddisplay(self, 'basin_id', 'basin number assigned to each element [unitless]'))
        s += '{}\n'.format(fielddisplay(self, 'subglacial_discharge', 'sum of subglacial discharge for each basin [m/d]'))
        s += '{}\n'.format(fielddisplay(self, 'beta0', 'basin-specific intercept values [°C] (if beta_1==0 mean=beta_0/(1-sum(phi)))'))
        s += '{}\n'.format(fielddisplay(self, 'beta1', 'basin-specific trend values [°C yr^(-1)]'))
        s += '{}\n'.format(fielddisplay(self, 'ar_order', 'order of the autoregressive model [unitless]'))
        s += '{}\n'.format(fielddisplay(self, 'ar_initialtime', 'initial time assumed in the autoregressive model parameterization [yr]'))
        s += '{}\n'.format(fielddisplay(self, 'ar_timestep', 'time resolution of the autoregressive model [yr]'))
        s += '{}\n'.format(fielddisplay(self, 'phi', 'basin-specific vectors of lag coefficients [unitless]'))
        return s
    #}}}

    def setdefaultparameters(self):  # {{{
        self.basin_id = np.nan
        self.num_basins = 0
        self.subglacial_discharge = np.nan
        self.ar_order = 0.0  # Autoregression model of order 0
        return self
    #}}}

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

        md = checkfield(md, 'fieldname', 'frontalforcings.num_basins', 'numel', 1, 'NaN', 1, 'Inf', 1, '>', 0)
        md = checkfield(md, 'fieldname', 'frontalforcings.basin_id', 'Inf', 1, '>=', 0, '<=', md.frontalforcings.num_basins, 'size', [md.mesh.numberofelements])
        md = checkfield(md, 'fieldname', 'frontalforcings.subglacial_discharge', '>=', 0, 'NaN', 1, 'Inf', 1, 'timeseries', 1)
        if len(np.shape(self.beta0)) == 1:
            self.beta0 = np.array([self.beta0])
            self.beta1 = np.array([self.beta1])
        md = checkfield(md, 'fieldname', 'frontalforcings.beta0', 'NaN', 1, 'Inf', 1, 'size', [1, md.frontalforcings.num_basins], 'numel', md.frontalforcings.num_basins)
        md = checkfield(md, 'fieldname', 'frontalforcings.beta1', 'NaN', 1, 'Inf', 1, 'size', [1, md.frontalforcings.num_basins], 'numel', md.frontalforcings.num_basins)
        md = checkfield(md, 'fieldname', 'frontalforcings.ar_order', 'numel', 1, 'NaN', 1, 'Inf', 1, '>=', 0)
        md = checkfield(md, 'fieldname', 'frontalforcings.ar_initialtime', 'numel', 1, 'NaN', 1, 'Inf', 1)
        md = checkfield(md, 'fieldname', 'frontalforcings.ar_timestep', 'numel', 1, 'NaN', 1, 'Inf', 1, '>=', md.timestepping.time_step) # Autoregression time step cannot be finer than ISSM timestep
        md = checkfield(md, 'fieldname', 'frontalforcings.phi', 'NaN', 1, 'Inf', 1, 'size', [md.frontalforcings.num_basins, md.frontalforcings.ar_order])
        return md
    # }}}

    def extrude(self, md):  # {{{
        # Nothing for now
        return self
    # }}}

    def marshall(self, prefix, md, fid):  # {{{
        yts = md.constants.yts
        WriteData(fid, prefix, 'name', 'md.frontalforcings.parameterization', 'data', 3, 'format', 'Integer')
        WriteData(fid, prefix, 'object', self, 'class', 'frontalforcings', 'fieldname', 'num_basins', 'format', 'Integer')
        WriteData(fid, prefix, 'object', self, 'class', 'frontalforcings', 'fieldname', 'subglacial_discharge', 'format', 'DoubleMat', 'mattype', 1, 'timeserieslength', md.mesh.numberofvertices + 1, 'yts', yts)
        WriteData(fid, prefix, 'object', self, 'class', 'frontalforcings', 'fieldname', 'ar_order', 'format', 'Integer')
        WriteData(fid, prefix, 'object', self, 'class', 'frontalforcings', 'fieldname', 'ar_initialtime', 'format', 'Double', 'scale', yts)
        WriteData(fid, prefix, 'object', self, 'class', 'frontalforcings', 'fieldname', 'ar_timestep', 'format', 'Double', 'scale', yts)
        WriteData(fid, prefix, 'object', self, 'class', 'frontalforcings', 'fieldname', 'basin_id', 'data', self.basin_id - 1, 'name', 'md.frontalforcings.basin_id', 'format', 'IntMat', 'mattype', 2)  # 0-indexed
        WriteData(fid, prefix, 'object', self, 'class', 'frontalforcings', 'fieldname', 'beta0', 'format', 'DoubleMat', 'name', 'md.frontalforcings.beta0')
        WriteData(fid, prefix, 'object', self, 'class', 'frontalforcings', 'fieldname', 'beta1', 'format', 'DoubleMat', 'name', 'md.frontalforcings.beta1', 'scale', 1 / yts, 'yts', yts)
        WriteData(fid, prefix, 'object', self, 'class', 'frontalforcings', 'fieldname', 'phi', 'format', 'DoubleMat', 'name', 'md.frontalforcings.phi', 'yts', yts)
    # }}}
