import numpy as np

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


class materials(object):
    '''
    MATERIALS class definition

       Usage:
          materials = materials()
    '''

    def __init__(self, *args): #{{{
        self.nature = []
        if not len(args):
            self.nature = ['ice']
        else:
            self.nature = args

        for i in range(len(self.nature)):
            if not(self.nature[i] == 'litho' or self.nature[i] == 'ice' or self.nature[i] == 'hydro'):
                raise RuntimeError("materials constructor error message: nature of the material not supported yet! ('ice' or 'litho' or 'hydro')")

        #start filling in the dynamic fields:
        for i in range(len(self.nature)):
            nat = self.nature[i]
            if nat == 'ice':
                setattr(self, 'rho_ice', 0)
                setattr(self, 'rho_water', 0)
                setattr(self, 'rho_freshwater', 0)
                setattr(self, 'mu_water', 0)
                setattr(self, 'heatcapacity', 0)
                setattr(self, 'latentheat', 0)
                setattr(self, 'thermalconductivity', 0)
                setattr(self, 'temperateiceconductivity', 0)
                setattr(self, 'meltingpoint', 0)
                setattr(self, 'beta', 0)
                setattr(self, 'mixed_layer_capacity', 0)
                setattr(self, 'thermal_exchange_velocity', 0)
                setattr(self, 'rheology_B', 0)
                setattr(self, 'rheology_n', 0)
                setattr(self, 'rheology_law', 0)
            elif nat == 'litho':
                setattr(self, 'numlayers', 0)
                setattr(self, 'radius', 0)
                setattr(self, 'viscosity', 0)
                setattr(self, 'lame_lambda', 0)
                setattr(self, 'lame_mu', 0)
                setattr(self, 'burgers_viscosity', 0)
                setattr(self, 'burgers_mu', 0)
                setattr(self, 'isburgers', 0)
                setattr(self, 'density', 0)
                setattr(self, 'issolid', 0)
            elif nat == 'hydro':
                setattr(self, 'rho_ice', 0)
                setattr(self, 'rho_water', 0)
                setattr(self, 'rho_freshwater', 0)
                setattr(self, 'earth_density', 0)
            else:
                raise RuntimeError("materials constructor error message: nature of the material not supported yet! ('ice' or 'litho' or 'hydro')")

        #set default parameters:
        self.setdefaultparameters()
    #}}}

    def setdefaultparameters(self): #{{{
        for i in range(len(self.nature)):
            nat = self.nature[i]
            if nat == 'ice':
                #ice density (kg/m^3)
                self.rho_ice = 917.
                #ocean water density (kg/m^3)
                self.rho_water = 1023.
                #fresh water density (kg/m^3)
                self.rho_freshwater = 1000.
                #water viscosity (N.s/m^2)
                self.mu_water = 0.001787
                #ice heat capacity cp (J/kg/K)
                self.heatcapacity = 2093.
                #ice latent heat of fusion L (J / kg)
                self.latentheat = 3.34e5
                #ice thermal conductivity (W/m/K)
                self.thermalconductivity = 2.4
                #wet ice thermal conductivity (W/m/K)
                self.temperateiceconductivity = 0.24
                #the melting point of ice at 1 atmosphere of pressure in K
                self.meltingpoint = 273.15
                #rate of change of melting point with pressure (K/Pa)
                self.beta = 9.8e-8
                #mixed layer (ice-water interface) heat capacity (J/kg/K)
                self.mixed_layer_capacity = 3974.
                #thermal exchange velocity (ice-water interface) (m/s)
                self.thermal_exchange_velocity = 1.00e-4
                #Rheology law: what is the temperature dependence of B with T
                #available: none, paterson and arrhenius
                self.rheology_law = 'Paterson'
            elif nat == 'litho':
                #we default to a configuration that enables running GIA solutions using giacaron and / or giaivins.
                self.numlayers = 2
                #center of the earth (approximation, must not be 0), then the lab (lithosphere / asthenosphere boundary) then the surface
                #(with 1d3 to avoid numerical singularities)
                self.radius = [1e3, 6278e3, 6378e3]
                self.viscosity = [1e21, 1e40]  #mantle and lithosphere viscosity (respectively) [Pa.s]
                self.lame_mu = [1.45e11, 6.7e10]  # (Pa)  #lithosphere and mantle shear modulus (respectively) [Pa]
                self.lame_lambda = self.lame_mu  # (Pa)  #mantle and lithosphere lamba parameter (respectively) [Pa]
                self.burgers_viscosity = [np.nan, np.nan]
                self.burgers_mu = [np.nan, np.nan]
                self.isburgers = [0, 0]
                self.density = [5.51e3, 5.50e3]  # (Pa)  #mantle and lithosphere density [kg/m^3]
                self.issolid = [1, 1]  # is layer solid or liquid.
            elif nat == 'hydro':
                #ice density (kg/m^3)
                self.rho_ice = 917.
                #ocean water density (kg/m^3)
                self.rho_water = 1023.
                #average density of the Earth (kg/m^3)
                self.earth_density = 5512
                #fresh water density (kg/m^3)
                self.rho_freshwater = 1000.
            else:
                raise RuntimeError("materials setdefaultparameters error message: nature of the material not supported yet! ('ice' or 'litho' or 'hydro')")
    #}}}

    def __repr__(self): #{{{
        string = "   Materials:"
        for i in range(len(self.nature)):
            nat = self.nature[i]
            if nat == 'ice':
                string = "%s\n%s" % (string, 'Ice:')
                string = "%s\n%s" % (string, fielddisplay(self, "rho_ice", "ice density [kg/m^3]"))
                string = "%s\n%s" % (string, fielddisplay(self, "rho_water", "ocean water density [kg/m^3]"))
                string = "%s\n%s" % (string, fielddisplay(self, "rho_freshwater", "fresh water density [kg/m^3]"))
                string = "%s\n%s" % (string, fielddisplay(self, "mu_water", "water viscosity [N s/m^2]"))
                string = "%s\n%s" % (string, fielddisplay(self, "heatcapacity", "heat capacity [J/kg/K]"))
                string = "%s\n%s" % (string, fielddisplay(self, "thermalconductivity", "ice thermal conductivity [W/m/K]"))
                string = "%s\n%s" % (string, fielddisplay(self, "temperateiceconductivity", "temperate ice thermal conductivity [W/m/K]"))
                string = "%s\n%s" % (string, fielddisplay(self, "meltingpoint", "melting point of ice at 1atm in K"))
                string = "%s\n%s" % (string, fielddisplay(self, "latentheat", "latent heat of fusion [J/m^3]"))
                string = "%s\n%s" % (string, fielddisplay(self, "beta", "rate of change of melting point with pressure [K/Pa]"))
                string = "%s\n%s" % (string, fielddisplay(self, "mixed_layer_capacity", "mixed layer capacity [W/kg/K]"))
                string = "%s\n%s" % (string, fielddisplay(self, "thermal_exchange_velocity", "thermal exchange velocity [m/s]"))
                string = "%s\n%s" % (string, fielddisplay(self, "rheology_B", "flow law parameter [Pa s^(1/n)]"))
                string = "%s\n%s" % (string, fielddisplay(self, "rheology_n", "Glen's flow law exponent"))
                string = "%s\n%s" % (string, fielddisplay(self, "rheology_law", "law for the temperature dependance of the rheology: 'None', 'BuddJacka', 'Cuffey', 'CuffeyTemperate', 'Paterson', 'Arrhenius', 'LliboutryDuval', 'NyeCO2', or 'NyeH2O'"))
            elif nat == 'litho':
                string = "%s\n%s" % (string, 'Litho:')
                string = "%s\n%s" % (string, fielddisplay(self, 'numlayers', 'number of layers (default 2)'))
                string = "%s\n%s" % (string, fielddisplay(self, 'radius', 'array describing the radius for each interface (numlayers + 1) [m]'))
                string = "%s\n%s" % (string, fielddisplay(self, 'viscosity', 'array describing each layer''s viscosity (numlayers) [Pa.s]'))
                string = "%s\n%s" % (string, fielddisplay(self, 'lame_lambda', 'array describing the lame lambda parameter (numlayers) [Pa]'))
                string = "%s\n%s" % (string, fielddisplay(self, 'lame_mu', 'array describing the shear modulus for each layers (numlayers) [Pa]'))
                string = "%s\n%s" % (string, fielddisplay(self, 'burgers_viscosity', 'array describing each layer''s transient viscosity, only for Burgers rheologies  (numlayers) [Pa.s]'))
                string = "%s\n%s" % (string, fielddisplay(self, 'burgers_mu', 'array describing each layer''s transient shear modulus, only for Burgers rheologies  (numlayers) [Pa]'))
                string = "%s\n%s" % (string, fielddisplay(self, 'isburgers', 'array describing whether we adopt a MaxWell (0) or Burgers (1) rheology (default 0)'))
                string = "%s\n%s" % (string, fielddisplay(self, 'density', 'array describing each layer''s density (numlayers) [kg/m^3]'))
                string = "%s\n%s" % (string, fielddisplay(self, 'issolid', 'array describing whether the layer is solid or liquid (default 1) (numlayers)'))
            elif nat == 'hydro':
                string = "%s\n%s" % (string, 'Hydro:')
                string = "%s\n%s" % (string, fielddisplay(self, "rho_ice", "ice density [kg/m^3]"))
                string = "%s\n%s" % (string, fielddisplay(self, "rho_water", "ocean water density [kg/m^3]"))
                string = "%s\n%s" % (string, fielddisplay(self, "earth_density", "mantle density [kg/m^3]"))
                string = "%s\n%s" % (string, fielddisplay(self, "rho_freshwater", "fresh water density [kg/m^3]"))

            else:
                raise RuntimeError("materials constructor error message: nature of the material not supported yet! ('ice' or 'litho' or 'hydro')")

        return string
    #}}}

    def checkconsistency(self, md, solution, analyses): #{{{
        for i in range(len(self.nature)):
            nat = self.nature[i]
            if nat == 'ice':
                md = checkfield(md, 'fieldname', 'materials.rho_ice', '>', 0)
                md = checkfield(md, 'fieldname', 'materials.rho_water', '>', 0)
                md = checkfield(md, 'fieldname', 'materials.rho_freshwater', '>', 0)
                md = checkfield(md, 'fieldname', 'materials.mu_water', '>', 0)
                md = checkfield(md, 'fieldname', 'materials.rheology_B', '>', 0, 'timeseries', 1, 'NaN', 1, 'Inf', 1)
                md = checkfield(md, 'fieldname', 'materials.rheology_n', '>', 0, 'size', [md.mesh.numberofelements])
                md = checkfield(md, 'fieldname', 'materials.rheology_law', 'values', ['None', 'BuddJacka', 'Cuffey', 'CuffeyTemperate', 'Paterson', 'Arrhenius', 'LliboutryDuval', 'NyeCO2', 'NyeH2O'])
            elif nat == 'litho':
                if 'LoveAnalysis' not in analyses:
                    return md

                md = checkfield(md, 'fieldname', 'materials.numlayers', 'NaN', 1, 'Inf', 1, '>', 0, 'numel', [1])
                md = checkfield(md, 'fieldname', 'materials.radius', 'NaN', 1, 'Inf', 1, 'size', [md.materials.numlayers + 1, 1], '>', 0)
                md = checkfield(md, 'fieldname', 'materials.lame_mu', 'NaN', 1, 'Inf', 1, 'size', [md.materials.numlayers, 1], '>=', 0)
                md = checkfield(md, 'fieldname', 'materials.lame_lambda', 'NaN', 1, 'Inf', 1, 'size', [md.materials.numlayers, 1], '>=', 0)
                md = checkfield(md, 'fieldname', 'materials.issolid', 'NaN', 1, 'Inf', 1, 'size', [md.materials.numlayers, 1], '>=', 0, '<', 2)
                md = checkfield(md, 'fieldname', 'materials.density', 'NaN', 1, 'Inf', 1, 'size', [md.materials.numlayers, 1], '>', 0)
                md = checkfield(md, 'fieldname', 'materials.viscosity', 'NaN', 1, 'Inf', 1, 'size', [md.materials.numlayers, 1], '>=', 0)
                md = checkfield(md, 'fieldname', 'materials.isburgers', 'NaN', 1, 'Inf', 1, 'size', [md.materials.numlayers, 1], '>=', 0, '<=', 1)
                md = checkfield(md, 'fieldname', 'materials.burgers_viscosity', 'Inf', 1, 'size', [md.materials.numlayers, 1], '>=', 0)
                md = checkfield(md, 'fieldname', 'materials.burgers_mu', 'Inf', 1, 'size', [md.materials.numlayers, 1], '>=', 0)

                for i in range(md.materials.numlayers):
                    if md.materials.isburgers[i] and (np.isnan(md.materials.burgers_viscosity[i] or np.isnan(md.materials.burgers_mu[i]))):
                        raise RuntimeError("materials checkconsistency error message: Litho burgers_viscosity or burgers_mu has NaN values, inconsistent with isburgers choice")

                    if md.materials.issolid[0] == 0 or md.materials.lame_mu[0] == 0:
                        raise RuntimeError('First layer must be solid (issolid(1) > 0 AND lame_mu(1) > 0). Add a weak inner core if necessary.')

                    for i in range(md.materials.numlayers - 1):
                        if (not md.materials.issolid[i]) and (not md.materials.issolid[i + 1]):  #if there are at least two consecutive indices that contain issolid = 0
                            raise RuntimeError("%s%i%s" % ('2 or more adjacent fluid layers detected starting at layer ', i, '. This is not supported yet. Consider merging them.'))
            elif nat == 'hydro':
                md = checkfield(md, 'fieldname', 'materials.rho_ice', '>', 0)
                md = checkfield(md, 'fieldname', 'materials.rho_water', '>', 0)
                md = checkfield(md, 'fieldname', 'materials.earth_density', '>', 0, 'numel', [1])
                md = checkfield(md, 'fieldname', 'materials.rho_freshwater', '>', 0)
            else:
                raise RuntimeError("materials checkconsistency error message: nature of the material not supported yet! ('ice' or 'litho' or 'hydro')")

        return md
    #}}}

    def marshall(self, prefix, md, fid): #{{{
        #1: MatdamageiceEnum 2: MatestarEnum 3: MaticeEnum 4: MatenhancediceEnum 5: MaterialsEnum
        WriteData(fid, prefix, 'name', 'md.materials.nature', 'data', naturetointeger(self.nature), 'format', 'IntMat', 'mattype', 3)
        WriteData(fid, prefix, 'name', 'md.materials.type', 'data', 5, 'format', 'Integer') #DANGER, this can evolve if you have classes

        for i in range(len(self.nature)):
            nat = self.nature[i]
            if nat == 'ice':
                WriteData(fid, prefix, 'name', 'md.materials.type', 'data', 3, 'format', 'Integer')
                WriteData(fid, prefix, 'object', self, 'class', 'materials', 'fieldname', 'rho_ice', 'format', 'Double')
                WriteData(fid, prefix, 'object', self, 'class', 'materials', 'fieldname', 'rho_water', 'format', 'Double')
                WriteData(fid, prefix, 'object', self, 'class', 'materials', 'fieldname', 'rho_freshwater', 'format', 'Double')
                WriteData(fid, prefix, 'object', self, 'class', 'materials', 'fieldname', 'mu_water', 'format', 'Double')
                WriteData(fid, prefix, 'object', self, 'class', 'materials', 'fieldname', 'heatcapacity', 'format', 'Double')
                WriteData(fid, prefix, 'object', self, 'class', 'materials', 'fieldname', 'latentheat', 'format', 'Double')
                WriteData(fid, prefix, 'object', self, 'class', 'materials', 'fieldname', 'thermalconductivity', 'format', 'Double')
                WriteData(fid, prefix, 'object', self, 'class', 'materials', 'fieldname', 'temperateiceconductivity', 'format', 'Double')
                WriteData(fid, prefix, 'object', self, 'class', 'materials', 'fieldname', 'meltingpoint', 'format', 'Double')
                WriteData(fid, prefix, 'object', self, 'class', 'materials', 'fieldname', 'beta', 'format', 'Double')
                WriteData(fid, prefix, 'object', self, 'class', 'materials', 'fieldname', 'mixed_layer_capacity', 'format', 'Double')
                WriteData(fid, prefix, 'object', self, 'class', 'materials', 'fieldname', 'thermal_exchange_velocity', 'format', 'Double')
                WriteData(fid, prefix, 'object', self, 'class', 'materials', 'fieldname', 'rheology_B', 'format', 'DoubleMat', 'mattype', 1, 'timeserieslength', md.mesh.numberofvertices + 1, 'yts', md.constants.yts)
                WriteData(fid, prefix, 'object', self, 'class', 'materials', 'fieldname', 'rheology_n', 'format', 'DoubleMat', 'mattype', 2)
                WriteData(fid, prefix, 'data', self.rheology_law, 'name', 'md.materials.rheology_law', 'format', 'String')
            elif nat == 'litho':
                WriteData(fid, prefix, 'object', self, 'class', 'materials', 'fieldname', 'numlayers', 'format', 'Integer')
                WriteData(fid, prefix, 'object', self, 'class', 'materials', 'fieldname', 'radius', 'format', 'DoubleMat', 'mattype', 3)
                WriteData(fid, prefix, 'object', self, 'class', 'materials', 'fieldname', 'lame_mu', 'format', 'DoubleMat', 'mattype', 3)
                WriteData(fid, prefix, 'object', self, 'class', 'materials', 'fieldname', 'lame_lambda', 'format', 'DoubleMat', 'mattype', 3)
                WriteData(fid, prefix, 'object', self, 'class', 'materials', 'fieldname', 'issolid', 'format', 'DoubleMat', 'mattype', 3)
                WriteData(fid, prefix, 'object', self, 'class', 'materials', 'fieldname', 'density', 'format', 'DoubleMat', 'mattype', 3)
                WriteData(fid, prefix, 'object', self, 'class', 'materials', 'fieldname', 'viscosity', 'format', 'DoubleMat', 'mattype', 3)
                WriteData(fid, prefix, 'object', self, 'class', 'materials', 'fieldname', 'isburgers', 'format', 'DoubleMat', 'mattype', 3)
                WriteData(fid, prefix, 'object', self, 'class', 'materials', 'fieldname', 'burgers_viscosity', 'format', 'DoubleMat', 'mattype', 3)
                WriteData(fid, prefix, 'object', self, 'class', 'materials', 'fieldname', 'burgers_mu', 'format', 'DoubleMat', 'mattype', 3)
            elif nat == 'hydro':
                WriteData(fid, prefix, 'object', self, 'class', 'materials', 'fieldname', 'rho_ice', 'format', 'Double')
                WriteData(fid, prefix, 'object', self, 'class', 'materials', 'fieldname', 'rho_water', 'format', 'Double')
                WriteData(fid, prefix, 'object', self, 'class', 'materials', 'fieldname', 'earth_density', 'format', 'Double')
                WriteData(fid, prefix, 'object', self, 'class', 'materials', 'fieldname', 'rho_freshwater', 'format', 'Double')
            else:
                raise RuntimeError("materials constructor error message: nature of the material not supported yet! ('ice' or 'litho' or 'hydro')")
    #}}}

    def extrude(self, md): #{{{
        for i in range(len(self.nature)):
            nat = self.nature[i]
            if nat == 'ice':
                self.rheology_B = project3d(md, 'vector', self.rheology_B, 'type', 'node')
                self.rheology_n = project3d(md, 'vector', self.rheology_n, 'type', 'element')
            return self
    #}}}
#}}}

def naturetointeger(strnat): #{{{
    intnat = np.zeros(len(strnat))

    for i in range(len(intnat)):
        if strnat[i] == 'damageice':
            intnat[i] = 1
        elif strnat[i] == 'estar':
            intnat[i] = 2
        elif strnat[i] == 'ice':
            intnat[i] = 3
        elif strnat[i] == 'enhancedice':
            intnat[i] = 4
        # elif strnat[i] == 'materials':
        #     intnat[i] = 5 #this case will never happen, kept to ensure equivalent of codes between IoCodeToMeterialsEnum and IoCodeToNatureEnum
        elif strnat[i] == 'litho':
            intnat[i] = 6
        elif strnat[i] == 'hydro':
            intnat[i] = 7
        else:
            raise RuntimeError("materials constructor error message: nature of the material not supported yet! ('ice' or 'litho' or 'hydro')")

    return intnat
#}}}
