# -*- coding: utf-8 -*-
"""

"""
from numpy import array, zeros, dot, ndarray, average

from pyfem.elements.BaseElement import BaseElement
from pyfem.fem.Timer import Timer
from pyfem.fem.constants import DTYPE
from pyfem.io.Dof import Dof
from pyfem.io.Material import Material
from pyfem.io.Section import Section
from pyfem.isoelements.IsoElementShape import IsoElementShape
from pyfem.materials.get_material_data import MaterialData
from pyfem.utils.colors import error_style


class Thermal(BaseElement):
    """
    温度单元。

    :ivar qp_temperatures: 积分点处的温度列表
    :vartype qp_temperatures: ndarray

    :ivar qp_heat_fluxes: 积分点处的热流密度列表
    :vartype qp_heat_fluxes: ndarray

    :ivar qp_ddsddts: 积分点处的材料热传导系数矩阵列表
    :vartype qp_ddsddts: list[ndarray]

    :ivar ntens: 总应力数量
    :vartype ntens: int

    :ivar ndi: 轴向应力数量
    :vartype ndi: int

    :ivar nshr: 剪切应力数量
    :vartype nshr: int
    """

    __slots_dict__: dict = {
        'qp_temperatures': ('ndarray', '积分点处的温度列表'),
        'qp_heat_fluxes': ('ndarray', '积分点处的热流密度列表'),
        'qp_ddsddts': ('list[ndarray]', '积分点处的材料热传导系数矩阵列表'),
        'ntens': ('int', '总应力数量'),
        'ndi': ('int', '轴向应力数量'),
        'nshr': ('int', '剪切应力数量')
    }

    __slots__: list = BaseElement.__slots__ + [slot for slot in __slots_dict__.keys()]

    def __init__(self, element_id: int,
                 iso_element_shape: IsoElementShape,
                 connectivity: ndarray,
                 node_coords: ndarray,
                 dof: Dof,
                 materials: list[Material],
                 section: Section,
                 material_data_list: list[MaterialData],
                 timer: Timer) -> None:

        super().__init__(element_id, iso_element_shape, connectivity, node_coords)

        self.allowed_material_data_list = [('ThermalIsotropic',)]
        self.allowed_material_number = 1

        self.dof = dof
        self.materials = materials
        self.section = section
        self.material_data_list = material_data_list
        self.check_materials()
        self.timer = timer

        self.dof_names = ['T']
        self.ntens = 6
        self.ndi = 3
        self.nshr = 3

        if dof.names != self.dof_names:
            error_msg = f'{dof.names} is not the supported dof of {type(self).__name__} element'
            raise NotImplementedError(error_style(error_msg))

        element_dof_number = len(self.dof_names) * self.iso_element_shape.nodes_number

        self.element_dof_number = element_dof_number
        self.element_dof_values = zeros(element_dof_number, dtype=DTYPE)
        self.element_ddof_values = zeros(element_dof_number, dtype=DTYPE)
        self.element_fint = zeros(element_dof_number, dtype=DTYPE)
        self.element_stiffness = zeros(shape=(self.element_dof_number, self.element_dof_number), dtype=DTYPE)

        self.qp_temperatures: list[ndarray] = None  # type: ignore
        self.qp_heat_fluxes: list[ndarray] = None  # type: ignore
        self.qp_ddsddts: list[ndarray] = None  # type: ignore

    def update_element_material_stiffness_fint(self,
                                               is_update_material: bool = True,
                                               is_update_stiffness: bool = True,
                                               is_update_fint: bool = True, ) -> None:
        element_id = self.element_id
        timer = self.timer
        ntens = self.ntens
        ndi = self.ndi
        nshr = self.nshr

        qp_number = self.qp_number
        qp_weight_times_jacobi_dets = self.qp_weight_times_jacobi_dets
        qp_shape_values = self.iso_element_shape.qp_shape_values
        qp_shape_gradients = self.iso_element_shape.qp_shape_gradients

        qp_state_variables = self.qp_state_variables
        qp_state_variables_new = self.qp_state_variables_new

        element_dof_values = self.element_dof_values
        element_ddof_values = self.element_ddof_values

        material_data = self.material_data_list[0]

        if is_update_stiffness:
            self.element_stiffness = zeros(shape=(self.element_dof_number, self.element_dof_number), dtype=DTYPE)

        if is_update_fint:
            self.element_fint = zeros(self.element_dof_number, dtype=DTYPE)

        if is_update_material:
            self.qp_ddsddts = list()
            self.qp_temperatures = list()
            self.qp_heat_fluxes = list()

        for i in range(qp_number):
            if is_update_material:
                qp_weight_times_jacobi_det = qp_weight_times_jacobi_dets[i]
                qp_shape_value = qp_shape_values[i]
                qp_shape_gradient = qp_shape_gradients[i]
                qp_temperature = dot(qp_shape_value, element_dof_values)
                qp_dtemperature = dot(qp_shape_value, element_ddof_values)
                qp_temperature_gradient = dot(qp_shape_gradient, element_dof_values)
                qp_dtemperature_gradient = dot(qp_shape_gradient, element_ddof_values)

                variable = {'temperature': qp_temperature,
                            'dtemperature': qp_dtemperature,
                            'temperature_gradient': qp_temperature_gradient,
                            'dtemperature_gradient': qp_dtemperature_gradient}
                qp_ddsddt, qp_output = material_data.get_tangent(variable=variable,
                                                                 state_variable=qp_state_variables[i],
                                                                 state_variable_new=qp_state_variables_new[i],
                                                                 element_id=element_id,
                                                                 iqp=i,
                                                                 ntens=ntens,
                                                                 ndi=ndi,
                                                                 nshr=nshr,
                                                                 timer=timer)
                qp_heat_flux = qp_output['heat_flux']
                self.qp_ddsddts.append(qp_ddsddt)
                self.qp_temperatures.append(qp_temperature)
                self.qp_heat_fluxes.append(qp_heat_flux)
            else:
                qp_shape_gradient = qp_shape_gradients[i]
                qp_weight_times_jacobi_det = qp_weight_times_jacobi_dets[i]
                qp_ddsddt = self.qp_ddsddts[i]
                qp_heat_flux = self.qp_heat_fluxes[i]

            if is_update_stiffness:
                self.element_stiffness += dot(qp_shape_gradient.transpose(), dot(qp_ddsddt, qp_shape_gradient)) * \
                                          qp_weight_times_jacobi_det

            if is_update_fint:
                self.element_fint += dot(qp_shape_gradient.transpose(), qp_heat_flux) * qp_weight_times_jacobi_det

    def update_element_field_variables(self) -> None:
        qp_temperatures = self.qp_temperatures
        qp_heat_fluxes = self.qp_heat_fluxes

        average_temperatures = average(qp_temperatures, axis=0)
        average_heat_fluxes = average(qp_heat_fluxes, axis=0)

        self.qp_field_variables['temperature'] = array(qp_temperatures, dtype=DTYPE)
        self.qp_field_variables['heat_flux'] = array(qp_heat_fluxes, dtype=DTYPE)

        self.element_average_field_variables['T'] = average_temperatures
        if len(average_heat_fluxes) >= 1:
            self.element_average_field_variables['HFL1'] = average_heat_fluxes[0]
        if len(average_heat_fluxes) >= 2:
            self.element_average_field_variables['HFL2'] = average_heat_fluxes[1]
        if len(average_heat_fluxes) >= 3:
            self.element_average_field_variables['HFL3'] = average_heat_fluxes[2]


if __name__ == "__main__":
    from pyfem.utils.visualization import print_slots_dict

    print_slots_dict(Thermal.__slots_dict__)

    from pyfem.Job import Job

    job = Job(r'F:\Github\pyfem\examples\thermal\1element\hex8\Job-1.toml')

    job.run()
