import mgis.behaviour as mgis_bv
import dolfin

mgis_hypothesis = {"plane_strain": mgis_bv.Hypothesis.PlaneStrain,
                   "plane_stress": mgis_bv.Hypothesis.PlaneStress,
                   "3d": mgis_bv.Hypothesis.Tridimensional,
                   "axisymmetric": mgis_bv.Hypothesis.Axisymmetrical}


class MFrontNonlinearMaterial:
    def __init__(self, path, name, hypothesis="3d",
                 material_properties={}, external_state_variables={"Temperature": 293.15}):
        self.path = path
        self.name = name
        # Defining the modelling hypothesis
        self.hypothesis = mgis_hypothesis[hypothesis]
        # if self.hypothesis == mgis_bv.Hypothesis.Tridimensional:
        #     self._tensorsize = 9
        # else:
        #     self._tensorsize = 5
        self.material_properties = material_properties
        self.external_state_variables = external_state_variables
        # finite strain options
        bopts = mgis_bv.FiniteStrainBehaviourOptions()
        bopts.stress_measure = mgis_bv.FiniteStrainBehaviourOptionsStressMeasure.PK1
        bopts.tangent_operator = mgis_bv.FiniteStrainBehaviourOptionsTangentOperator.DPK1_DF
        # Loading the behaviour
        self.behaviour = mgis_bv.load(self.path, self.name, self.hypothesis)
        self.finite_strain = self.behaviour.getBehaviourType()=="StandardFiniteStrainBehaviour"
        if self.finite_strain:
            self.behaviour = mgis_bv.load(bopts, self.path, self.name, self.hypothesis)

    def set_data_manager(self, ngauss):
        # Setting the material data manager
        self.data_manager = mgis_bv.MaterialDataManager(self.behaviour, ngauss)
        self.update_material_properties()
        self.update_external_state_variable()

    def update_material_properties(self, material_properties=None):
        if material_properties is not None:
            self.material_properties = material_properties
        for s in [self.data_manager.s0, self.data_manager.s1]:
            for (key, value) in self.material_properties.items():
                if type(value) in [int, float]:
                    mgis_bv.setMaterialProperty(s, key, value)
                else:
                    if isinstance(value, dolfin.Function):
                        value = value.vector().get_local()
                    mgis_bv.setMaterialProperty(s, key, value, mgis_bv.MaterialStateManagerStorageMode.LocalStorage)


    def update_external_state_variable(self, external_state_variables=None):
        if external_state_variables is not None:
            self.external_state_variables = external_state_variables
        for s in [self.data_manager.s0, self.data_manager.s1]:
            for (key, value) in self.external_state_variables.items():
                mgis_bv.setExternalStateVariable(s, key, value)

    def get_state_variable_names(self):
        return [svar.name for svar in self.behaviour.internal_state_variables]
    def get_gradient_names(self):
        return [svar.name for svar in self.behaviour.gradients]
    def get_flux_names(self):
        return [svar.name for svar in self.behaviour.thermodynamic_forces]
    def get_state_variable_sizes(self):
        return [mgis_bv.getVariableSize(svar, self.hypothesis) for svar in self.behaviour.internal_state_variables]
    def get_gradient_sizes(self):
        return [mgis_bv.getVariableSize(svar, self.hypothesis) for svar in self.behaviour.gradients]
    def get_flux_sizes(self):
        return [mgis_bv.getVariableSize(svar, self.hypothesis) for svar in self.behaviour.thermodynamic_forces]
    def get_tangent_block_names(self):
        return [(t[0].name, t[1].name) for t in self.behaviour.tangent_operator_blocks]
    def get_tangent_block_sizes(self):
        return [tuple([mgis_bv.getVariableSize(tt, self.hypothesis) for tt in t]) \
                for t in self.behaviour.tangent_operator_blocks]