import numpy as np
from fielddisplay import fielddisplay
from checkfield import checkfield
from WriteData import WriteData
from project3d import project3d
from MatlabFuncs import *
from helpers import *


class SMBpddSicopolis(object):
    """
    SMBpddSicopolis Class definition

    Usage:
        SMBpddSicopolis = SMBpddSicopolis()
"""

    def __init__(self):  # {{{
        self.precipitation = float('NaN')
        self.monthlytemperatures = float('NaN')
        self.temperature_anomaly = float('NaN')
        self.precipitation_anomaly = float('NaN')
        self.smb_corr = float('NaN')
        self.desfac = 0
        self.s0p = float('NaN')
        self.s0t = float('NaN')
        self.rlaps = 0
        self.isfirnwarming = 0
        self.steps_per_step = 1
        self.requested_outputs = []

        self.setdefaultparameters()
    # }}}

    @staticmethod
    def SMBpddSicopolis(*args):  # {{{
        nargin = len(args)

        if nargin == 0:
            return SMBpddSicopolis()
        else:
            raise RuntimeError('SMBpddSicopolis: constructor not supported')
    # }}}

    def extrude(self, md):  # {{{
        self.precipitation = project3d(md, 'vector', self.precipitation, 'type', 'node')
        self.monthlytemperatures = project3d(md, 'vector', self.monthlytemperatures, 'type', 'node')
        self.temperature_anomaly = project3d(md, 'vector', self.temperature_anomaly, 'type', 'node')
        self.precipitation_anomaly = project3d(md, 'vector', self.precipitation_anomaly, 'type', 'node')
        self.smb_corr = project3d(md, 'vector', self.smb_corr, 'type', 'node')
        self.s0p = project3d(md, 'vector', self.s0p, 'type', 'node')
        self.s0t = project3d(md, 'vector', self.s0t, 'type', 'node')
    # }}}

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

    def initialize(self, md):  # {{{

        if np.isnan(self.s0p):
            self.s0p = np.zeros((md.mesh.numberofvertices, ))
            print('      no SMBpddSicopolis.s0p specified: values set as zero')

        if np.isnan(self.s0t):
            self.s0t = np.zeros((md.mesh.numberofvertices, ))
            print('      no SMBpddSicopolis.s0t specified: values set as zero')

        if np.isnan(self.temperature_anomaly):
            self.temperature_anomaly = np.zeros((md.mesh.numberofvertices, ))
            print('      no SMBpddSicopolis.temperature_anomaly specified: values set as zero')

        if np.isnan(self.precipitation_anomaly):
            self.precipitation_anomaly = np.ones((md.mesh.numberofvertices, ))
            print('      no SMBpddSicopolis.precipitation_anomaly specified: values set as ones')

        if np.isnan(self.smb_corr):
            self.smb_corr = np.zeros((md.mesh.numberofvertices, ))
            print('      no SMBpddSicopolis.smb_corr specified: values set as zero')
    # }}}

    def setdefaultparameters(self):  # {{{
        self.isfirnwarming = 1
        self.desfac = -np.log(2.0) / 1000
        self.rlaps = 7.4

    # }}}

    def checkconsistency(self, md, solution, analyses):  # {{{
        if (strcmp(solution, 'TransientSolution') and md.transient.issmb == 0):
            return

        if 'MasstransportAnalysis' in analyses:
            md = checkfield(md, 'fieldname', 'smb.desfac', '<=', 1, 'numel', 1)
            md = checkfield(md, 'fieldname', 'smb.s0p', '>=', 0, 'NaN', 1, 'Inf', 1, 'size', [md.mesh.numberofvertices, 1])
            md = checkfield(md, 'fieldname', 'smb.s0t', '>=', 0, 'NaN', 1, 'Inf', 1, 'size', [md.mesh.numberofvertices, 1])
            md = checkfield(md, 'fieldname', 'smb.rlaps', '>=', 0, 'numel', 1)
            md = checkfield(md, 'fieldname', 'smb.monthlytemperatures', 'NaN', 1, 'Inf', 1, 'size', [md.mesh.numberofvertices, 12])
            md = checkfield(md, 'fieldname', 'smb.precipitation','NaN', 1, 'Inf', 1, 'size', [md.mesh.numberofvertices, 12])

        md = checkfield(md, 'fieldname', 'smb.steps_per_step', '>=', 1, 'numel', [1])
        md = checkfield(md, 'fieldname', 'smb.requested_outputs', 'stringrow', 1)

        return md
    # }}}

    def __repr__(self):  # {{{
        string = '   surface forcings parameters:'
        string += '\n   SICOPOLIS PDD scheme (Calov & Greve, 2005) :'

        string = "%s\n%s" % (string, fielddisplay(self, 'monthlytemperatures', 'monthly surface temperatures [K]'))
        string = "%s\n%s" % (string, fielddisplay(self, 'precipitation', 'monthly surface precipitation [m/yr water eq]'))
        string = "%s\n%s" % (string, fielddisplay(self, 'temperature_anomaly', 'anomaly to monthly reference temperature (additive [K])'))
        string = "%s\n%s" % (string, fielddisplay(self, 'precipitation_anomaly', 'anomaly to monthly precipitation (multiplicative, e.g. q = q0*exp(0.070458*DeltaT) after Huybrechts (2002)) [no unit])'))
        string = "%s\n%s" % (string, fielddisplay(self, 'smb_corr', 'correction of smb after PDD call [m/a]'))
        string = "%s\n%s" % (string, fielddisplay(self, 's0p', 'should be set to elevation from precip source (between 0 and a few 1000s m, default is 0) [m]'))
        string = "%s\n%s" % (string, fielddisplay(self, 's0t', 'should be set to elevation from temperature source (between 0 and a few 1000s m, default is 0) [m]'))
        string = "%s\n%s" % (string, fielddisplay(self, 'rlaps', 'present day lapse rate (default is 7.4 degree/km)'))
        string = "%s\n%s" % (string, fielddisplay(self, 'desfac', 'desertification elevation factor (default is -log(2.0)/1000)'))
        string = "%s\n%s" % (string, fielddisplay(self, 'isfirnwarming', 'is firnwarming (Reeh 1991) activated (0 or 1, default is 1)'))
        string = "%s\n%s" % (string, fielddisplay(self, 'steps_per_step', 'number of smb steps per time step'))
        string = "%s\n%s" % (string, fielddisplay(self, 'requested_outputs', 'additional outputs requested (TemperaturePDD, SmbAccumulation, SmbMelt)'))
    # }}}

    def marshall(self, prefix, md, fid):  # {{{

        yts = md.constants.yts

        WriteData(fid, prefix, 'name', 'md.smb.model', 'data', 10, 'format', 'Integer')

        WriteData(fid, prefix, 'object', self, 'class', 'smb', 'fieldname', 'isfirnwarming', 'format', 'Boolean')
        WriteData(fid, prefix, 'object', self, 'class', 'smb', 'fieldname', 'desfac', 'format', 'Double')
        WriteData(fid, prefix, 'object', self, 'class', 'smb', 'fieldname', 's0p', 'format', 'DoubleMat', 'mattype', 1)
        WriteData(fid, prefix, 'object', self, 'class', 'smb', 'fieldname', 's0t', 'format', 'DoubleMat', 'mattype', 1)
        WriteData(fid, prefix, 'object', self, 'class', 'smb', 'fieldname', 'rlaps', 'format', 'Double')

        WriteData(fid, prefix, 'object', self, 'class', 'smb', 'fieldname', 'monthlytemperatures', 'format', 'DoubleMat', 'mattype', 1, 'timeserieslength', md.mesh.numberofvertices + 1, 'yts', md.constants.yts)
        WriteData(fid, prefix, 'object', self, 'class', 'smb', 'fieldname', 'precipitation', 'format', 'DoubleMat', 'mattype', 1, 'scale', 1. / yts, 'timeserieslength', md.mesh.numberofvertices + 1, 'yts', yts)
        WriteData(fid, prefix, 'object', self, 'class', 'smb', 'fieldname', 'temperature_anomaly', 'format', 'DoubleMat', 'mattype', 1, 'timeserieslength', md.mesh.numberofvertices + 1, 'yts', yts)
        WriteData(fid, prefix, 'object', self, 'class', 'smb', 'fieldname', 'precipitation_anomaly', 'format', 'DoubleMat', 'mattype', 1, 'scale', 1. / yts, 'timeserieslength', md.mesh.numberofvertices + 1, 'yts', yts)
        WriteData(fid, prefix, 'object', self, 'class', 'smb', 'fieldname', 'smb_corr', 'format', 'DoubleMat', 'mattype', 1, 'scale', 1. / yts, 'timeserieslength', md.mesh.numberofvertices + 1, 'yts', yts)
        WriteData(fid, prefix, 'object', self, 'fieldname', 'steps_per_step', 'format', 'Integer')

        #process requested outputs
        outputs = self.requested_outputs
        pos = np.where('default' in outputs)
        if not isempty(pos):
            outputs[pos] = []  #remove 'default' from outputs
            outputs = [outputs, defaultoutputs(self, md)]  #add defaults

        WriteData(fid, prefix, 'data', outputs, 'name', 'md.smb.requested_outputs', 'format', 'StringArray')

    # }}}
