import numpy as np

from checkfield import checkfield
from fielddisplay import fielddisplay
from lovenumbers import lovenumbers
from planetradius import planetradius
from project3d import project3d
from rotational import rotational
from solidearthsettings import solidearthsettings
from surfaceload import surfaceload
from WriteData import WriteData


class solidearth(object):
    """SOLIDEARTH class definition

    Usage:
        solidearth = solidearth()
    """

    def __init__(self, *args):  #{{{
        self.initialsealevel    = np.nan
        self.settings           = solidearthsettings()
        self.external           = []
        self.surfaceload        = surfaceload()
        self.lovenumbers        = lovenumbers()
        self.rotational         = rotational()
        self.planetradius       = planetradius('earth')
        self.requested_outputs  = []
        self.transitions        = []
        self.partitionice       = []
        self.partitionhydro     = []

        nargs = len(args)

        if nargs == 0:
            self.setdefaultparameters('earth')
        elif nargs == 1:
            self.setdefaultparameters(args[0])
        else:
            raise Exception('solidearth constructor error message: zero or one argument only!')
    #}}}

    def __repr__(self):  # {{{
        s = '   solidearthinputs, forcings and settings:\n'
        s += '{}\n'.format(fielddisplay(self, 'initialsealevel', 'sea level at the start of computation [m]'))
        s += '{}\n'.format(fielddisplay(self, 'planetradius', 'planet radius [m]'))
        s += '{}\n'.format(fielddisplay(self, 'transitions', 'indices into parts of the mesh that will be icecaps'))
        s += '{}\n'.format(fielddisplay(self, 'requested_outputs', 'additional outputs requested'))
        s += '{}\n'.format(fielddisplay(self, 'partitionice', 'ice partition vector for barystatic contribution'))
        s += '{}\n'.format(fielddisplay(self, 'partitionhydro', 'hydro partition vector for barystatic contribution'))
        if not self.external:
            s += '{}\n'.format(fielddisplay(self, 'external', 'external solution, of the type solidearthsolution'))
        print(self.settings)
        print(self.surfaceload)
        print(self.lovenumbers)
        print(self.rotational)
        if self.external:
            print(self.external)
        return s
    #}}}

    def setdefaultparameters(self, planet):  # {{{
        # Default output
        self.requested_outputs = ['default']

        # Transitions should be a list
        self.transitions = []

        # No partitions requested for barystatic contribution
        self.partitionice = []
        self.partitionhydro = []

        # No external solutions by defalt
        self.external = []

        # Earth radius
        self.planetradius = planetradius(planet)
    #}}}

    def checkconsistency(self, md, solution, analyses):  # {{{
        if ('SealevelriseAnalysis' not in analyses) or (solution == 'TransientSolution' and not md.transient.isslr):
            return md
        md = checkfield(md, 'fieldname', 'solidearth.initialsealevel', 'NaN', 1, 'Inf', 1, 'size', [md.mesh.numberofvertices])
        md = checkfield(md, 'fieldname', 'solidearth.requested_outputs', 'stringrow', 1)

        self.settings.checkconsistency(md, solution, analyses)
        self.surfaceload.checkconsistency(md, solution, analyses)
        self.lovenumbers.checkconsistency(md, solution, analyses)
        self.rotational.checkconsistency(md, solution, analyses)
        if self.external:
            if not isinstance(self.external,'solidearthsolution'):
                raise Exception('solidearth consistency check: external field should be a solidearthsolution')
            end
            self.external.checkconsistency(md,solution,analyses)
        return md
    #}}}

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

    def marshall(self, prefix, md, fid):  #{{{
        WriteData(fid, prefix, 'object', self, 'fieldname', 'initialsealevel', 'mattype', 1, 'format', 'DoubleMat', 'timeserieslength', md.mesh.numberofvertices + 1, 'yts', md.constants.yts)
        WriteData(fid, prefix, 'object', self, 'fieldname', 'planetradius', 'format', 'Double')
        WriteData(fid, prefix, 'object', self, 'fieldname', 'transitions', 'format', 'MatArray')

        if len(self.partitionice):
            npartice = np.max(self.partitionice) + 2
        else:
            npartice = 0

        if len(self.partitionhydro):
            nparthydro = np.max(self.partitionhydro) + 2
        else:
            nparthydro = 0

        WriteData(fid, prefix, 'object', self, 'fieldname', 'partitionice', 'mattype', 1, 'format', 'DoubleMat');
        WriteData(fid, prefix, 'data', npartice, 'format', 'Integer', 'name', 'md.solidearth.npartice');
        WriteData(fid, prefix, 'object', self, 'fieldname', 'partitionhydro', 'mattype', 1, 'format', 'DoubleMat');
        WriteData(fid, prefix, 'data', nparthydro,'format', 'Integer', 'name','md.solidearth.nparthydro');

        self.settings.marshall(prefix, md, fid)
        self.surfaceload.marshall(prefix, md, fid)
        self.lovenumbers.marshall(prefix, md, fid)
        self.rotational.marshall(prefix, md, fid)
        if self.external:
            self.external.marshall(prefix, md, fid)

        #process requested outputs
        outputs = self.requested_outputs
        pos = np.where(np.asarray(outputs) == 'default')[0]
        if len(pos):
            outputs = np.delete(outputs, pos)  #remove 'default' from outputs
            outputs = np.append(outputs, self.defaultoutputs(md))  #add defaults
        WriteData(fid, prefix, 'data', outputs, 'name', 'md.solidearth.requested_outputs', 'format', 'StringArray')
    #}}}

    def extrude(self, md): #{{{
        self.initialsealevel = project3d(md, 'vector', self.initialsealevel, 'type', 'node')
        return self
    #}}}