import numpy as np

from checkfield import checkfield
from fielddisplay import fielddisplay
from project3d import project3d
from WriteData import WriteData


class initialization(object):
    """INITIALIZATION class definition

    Usage:
    initialization = initialization()
    """

    def __init__(self):  # {{{
        self.vx = np.nan
        self.vy = np.nan
        self.vz = np.nan
        self.vel = np.nan
        self.enthalpy = np.nan
        self.pressure = np.nan
        self.temperature = np.nan
        self.waterfraction = np.nan
        self.watercolumn = np.nan
        self.sediment_head = np.nan
        self.epl_head = np.nan
        self.epl_thickness = np.nan
        self.hydraulic_potential = np.nan
        self.channelarea = np.nan

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

    def __repr__(self):  # {{{
        s = '   initial field values:\n'
        s += '{}\n'.format(fielddisplay(self, 'vx', 'x component of velocity [m / yr]'))
        s += '{}\n'.format(fielddisplay(self, 'vy', 'y component of velocity [m / yr]'))
        s += '{}\n'.format(fielddisplay(self, 'vz', 'z component of velocity [m / yr]'))
        s += '{}\n'.format(fielddisplay(self, 'vel', 'velocity norm [m / yr]'))
        s += '{}\n'.format(fielddisplay(self, 'pressure', 'pressure [Pa]'))
        s += '{}\n'.format(fielddisplay(self, 'temperature', 'temperature [K]'))
        s += '{}\n'.format(fielddisplay(self, 'enthalpy', 'enthalpy [J]'))
        s += '{}\n'.format(fielddisplay(self, 'waterfraction', 'fraction of water in the ice'))
        s += '{}\n'.format(fielddisplay(self, 'watercolumn', 'thickness of subglacial water [m]'))
        s += '{}\n'.format(fielddisplay(self, 'sediment_head', 'sediment water head of subglacial system [m]'))
        s += '{}\n'.format(fielddisplay(self, 'epl_head', 'epl water head of subglacial system [m]'))
        s += '{}\n'.format(fielddisplay(self, 'epl_thickness', 'thickness of the epl [m]'))
        s += '{}\n'.format(fielddisplay(self, 'hydraulic_potential', 'Hydraulic potential (for GlaDS) [Pa]'))
        s += '{}\n'.format(fielddisplay(self, 'channelarea', 'subglaciale water channel area (for GlaDS) [m2]'))
        return s
    #}}}

    def extrude(self, md):  # {{{
        self.vx = project3d(md, 'vector', self.vx, 'type', 'node')
        self.vy = project3d(md, 'vector', self.vy, 'type', 'node')
        self.vz = project3d(md, 'vector', self.vz, 'type', 'node')
        self.vel = project3d(md, 'vector', self.vel, 'type', 'node')
        self.temperature = project3d(md, 'vector', self.temperature, 'type', 'node')
        self.enthalpy = project3d(md, 'vector', self.enthalpy, 'type', 'node')
        self.waterfraction = project3d(md, 'vector', self.waterfraction, 'type', 'node')
        self.watercolumn = project3d(md, 'vector', self.watercolumn, 'type', 'node')
        self.sediment_head = project3d(md, 'vector', self.sediment_head, 'type', 'node', 'layer', 1)
        self.epl_head = project3d(md, 'vector', self.epl_head, 'type', 'node', 'layer', 1)
        self.epl_thickness = project3d(md, 'vector', self.epl_thickness, 'type', 'node', 'layer', 1)

        #Lithostatic pressure by default
        if np.ndim(md.geometry.surface) == 2:
            print('Reshaping md.geometry.surface for your convenience but you should fix it in your model set up')
            self.pressure = md.constants.g * md.materials.rho_ice * (md.geometry.surface.reshape(-1, ) - md.mesh.z)
        else:
            self.pressure = md.constants.g * md.materials.rho_ice * (md.geometry.surface - md.mesh.z)

        return self
    #}}}

    def setdefaultparameters(self):  # {{{
        return self
    #}}}

    def checkconsistency(self, md, solution, analyses):  # {{{
        if 'StressbalanceAnalysis' in analyses and not solution == 'TransientSolution' and not md.transient.isstressbalance:
            if not np.any(np.logical_or(np.isnan(md.initialization.vx), np.isnan(md.initialization.vy))):
                md = checkfield(md, 'fieldname', 'initialization.vx', 'NaN', 1, 'Inf', 1, 'size', [md.mesh.numberofvertices])
                md = checkfield(md, 'fieldname', 'initialization.vy', 'NaN', 1, 'Inf', 1, 'size', [md.mesh.numberofvertices])
        if 'MasstransportAnalysis' in analyses and not solution == 'TransientSolution' and not md.transient.ismasstransport:
            md = checkfield(md, 'fieldname', 'initialization.vx', 'NaN', 1, 'Inf', 1, 'size', [md.mesh.numberofvertices])
            md = checkfield(md, 'fieldname', 'initialization.vy', 'NaN', 1, 'Inf', 1, 'size', [md.mesh.numberofvertices])
        if 'BalancethicknessAnalysis' in analyses and solution == 'BalancethicknessSolution':
            md = checkfield(md, 'fieldname', 'initialization.vx', 'NaN', 1, 'Inf', 1, 'size', [md.mesh.numberofvertices])
            md = checkfield(md, 'fieldname', 'initialization.vy', 'NaN', 1, 'Inf', 1, 'size', [md.mesh.numberofvertices])
            # Triangle with zero velocity
            if np.any(np.logical_and(np.sum(np.abs(md.initialization.vx[md.mesh.elements - 1]), axis=1) == 0,
                                     np.sum(np.abs(md.initialization.vy[md.mesh.elements - 1]), axis=1) == 0)):
                md.checkmessage("at least one triangle has all its vertices with a zero velocity")
        if 'ThermalAnalysis' in analyses and not solution == 'TransientSolution' and not md.transient.isthermal:
            md = checkfield(md, 'fieldname', 'initialization.vx', 'NaN', 1, 'Inf', 1, 'size', [md.mesh.numberofvertices])
            md = checkfield(md, 'fieldname', 'initialization.vy', 'NaN', 1, 'Inf', 1, 'size', [md.mesh.numberofvertices])
            md = checkfield(md, 'fieldname', 'initialization.temperature', 'NaN', 1, 'Inf', 1, 'size', 'universal')
            if md.mesh.dimension() == 3:
                md = checkfield(md, 'fieldname', 'initialization.vz', 'NaN', 1, 'Inf', 1, 'size', [md.mesh.numberofvertices])
            md = checkfield(md, 'fieldname', 'initialization.pressure', 'NaN', 1, 'Inf', 1, 'size', [md.mesh.numberofvertices])
        if 'EnthalpyAnalysis' in analyses and md.thermal.isenthalpy:
            md = checkfield(md, 'fieldname', 'initialization.waterfraction', '>=', 0, 'size', [md.mesh.numberofvertices])
            md = checkfield(md, 'fieldname', 'initialization.watercolumn', '>=', 0, 'size', [md.mesh.numberofvertices])
            pos = np.nonzero(md.initialization.waterfraction > 0.)[0]
            if(pos.size):
                md = checkfield(md, 'fieldname', 'delta Tpmp', 'field', np.absolute(md.initialization.temperature[pos] - (md.materials.meltingpoint - md.materials.beta * md.initialization.pressure[pos])), '<', 1e-11, 'message', 'set temperature to pressure melting point at locations with waterfraction > 0')
        if 'HydrologyShreveAnalysis' in analyses:
            if hasattr(md.hydrology, 'hydrologyshreve'):
                md = checkfield(md, 'fieldname', 'initialization.watercolumn', 'NaN', 1, 'Inf', 1, 'size', [md.mesh.numberofvertices])
        if 'HydrologyDCInefficientAnalysis' in analyses:
            if hasattr(md.hydrology, 'hydrologydc'):
                md = checkfield(md, 'fieldname', 'initialization.sediment_head', 'NaN', 1, 'Inf', 1, 'size', [md.mesh.numberofvertices])
        if 'HydrologyDCEfficientAnalysis' in analyses:
            if hasattr(md.hydrology, 'hydrologydc'):
                if md.hydrology.isefficientlayer == 1:
                    md = checkfield(md, 'fieldname', 'initialization.epl_head', 'NaN', 1, 'Inf', 1, 'size', [md.mesh.numberofvertices])
                    md = checkfield(md, 'fieldname', 'initialization.epl_thickness', 'NaN', 1, 'Inf', 1, 'size', [md.mesh.numberofvertices])
        if 'HydrologyGlaDSAnalysis' in analyses:
            if hasattr(md.hydrology, 'hydrologyglads'):
                md = checkfield(md, 'fieldname', 'initialization.watercolumn', 'NaN', 1, 'Inf', 1, 'size', [md.mesh.numberofvertices])
                md = checkfield(md, 'fieldname', 'initialization.hydraulic_potential', 'NaN', 1, 'Inf', 1, 'size', [md.mesh.numberofvertices])
                md = checkfield(md, 'fieldname', 'initialization.channelarea', 'NaN', 1, 'Inf', 1, '>=', 0, 'size', [md.mesh.numberofelements])
        return md
    # }}}

    def marshall(self, prefix, md, fid):  # {{{
        yts = md.constants.yts
        WriteData(fid, prefix, 'object', self, 'fieldname', 'vx', 'format', 'DoubleMat', 'mattype', 1, 'scale', 1. / yts)
        WriteData(fid, prefix, 'object', self, 'fieldname', 'vy', 'format', 'DoubleMat', 'mattype', 1, 'scale', 1. / yts)
        WriteData(fid, prefix, 'object', self, 'fieldname', 'vz', 'format', 'DoubleMat', 'mattype', 1, 'scale', 1. / yts)
        WriteData(fid, prefix, 'object', self, 'fieldname', 'pressure', 'format', 'DoubleMat', 'mattype', 1)
        WriteData(fid, prefix, 'object', self, 'fieldname', 'temperature', 'format', 'DoubleMat', 'mattype', 1)
        WriteData(fid, prefix, 'object', self, 'fieldname', 'waterfraction', 'format', 'DoubleMat', 'mattype', 1)
        WriteData(fid, prefix, 'object', self, 'fieldname', 'sediment_head', 'format', 'DoubleMat', 'mattype', 1)
        WriteData(fid, prefix, 'object', self, 'fieldname', 'epl_head', 'format', 'DoubleMat', 'mattype', 1)
        WriteData(fid, prefix, 'object', self, 'fieldname', 'epl_thickness', 'format', 'DoubleMat', 'mattype', 1)
        WriteData(fid, prefix, 'object', self, 'fieldname', 'watercolumn', 'format', 'DoubleMat', 'mattype', 1)
        WriteData(fid, prefix, 'object', self, 'fieldname', 'channelarea', 'format', 'DoubleMat', 'mattype', 1)
        WriteData(fid, prefix, 'object', self, 'fieldname', 'hydraulic_potential', 'format', 'DoubleMat', 'mattype', 1)
        if md.thermal.isenthalpy:
            if (np.size(self.enthalpy) <= 1):
                tpmp = md.materials.meltingpoint - md.materials.beta * md.initialization.pressure
                pos = np.nonzero(md.initialization.waterfraction > 0.)[0]
                self.enthalpy = md.materials.heatcapacity * (md.initialization.temperature - md.constants.referencetemperature)
                self.enthalpy[pos] = md.materials.heatcapacity * (tpmp[pos].reshape(-1, ) - md.constants.referencetemperature) + md.materials.latentheat * md.initialization.waterfraction[pos].reshape(-1, )

            WriteData(fid, prefix, 'data', self.enthalpy, 'format', 'DoubleMat', 'mattype', 1, 'name', 'md.initialization.enthalpy')
    # }}}
