# Copyright (C) 2023 - 2025 ANSYS, Inc. and/or its affiliates.
# SPDX-License-Identifier: MIT
#
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

from collections.abc import Sequence
from dataclasses import dataclass

from ansys.dpf.core import MeshedRegion, Operator

from ansys.dpf.composites.layup_info import (
    AnalysisPlyInfoProvider,
    LayupPropertiesProvider,
    get_all_analysis_ply_names,
)


@dataclass(frozen=True)
class SolidStack:
    """
    List of solid elements and associated plies which build a stack.

    The stack of solid elements forms a laminate. It can be used
    to perform for example through-the-thickness sampling and laminate
    specific failure analysis.
    """

    # List of solid element labels ordered from the bottom to the top
    element_ids: Sequence[int]
    # List of analysis plies for each solid element in the stack.
    # Indexed by element ID (label)
    element_wise_analysis_plies: dict[int, Sequence[str]]
    # Element-wise ply thicknesses. Indexed by element ID (label)
    element_ply_thicknesses: dict[int, Sequence[float]]
    # The level defines the position of each element in the stack.
    # Multiple homogeneous elements can have the same level.
    # Indexed by element ID (label)
    element_wise_levels: dict[int, int]

    @property
    def number_of_elements(self) -> int:
        """Number of solid elements in the stack."""
        return len(self.element_ids)

    @property
    def element_ids_per_level(self) -> dict[int, list[int]]:
        """
        Element labels for each level in the stack.

        Note that a level (position in the stack) can have multiple elements
        in case of drop-offs or cut-offs. These elements have the
        same ``artificial` plies but needs special handling.
        """
        levels = {level for e, level in self.element_wise_levels.items()}
        element_ids_per_level: dict[int, list[int]] = {l: [] for l in levels}
        for e, level in self.element_wise_levels.items():
            element_ids_per_level[level].append(e)
        return element_ids_per_level

    @property
    def analysis_ply_ids_and_thicknesses(self) -> Sequence[tuple[str, float]]:
        """List of analysis plies and their thicknesses from the bottom to the top."""
        res = []
        for _, element_ids in self.element_ids_per_level.items():
            # Only the first element is of interest since all elements
            # on a level have the same plies
            e_id = element_ids[0]
            for index, ply_id in enumerate(self.element_wise_analysis_plies[e_id]):
                res.append((ply_id, self.element_ply_thicknesses[e_id][index]))

        return res

    @property
    def number_of_analysis_plies(self) -> int:
        """Number of analysis plies in the stack."""
        return len(self.analysis_ply_ids_and_thicknesses)


class SolidStackProvider:
    """
    Helper to get solid stack information from a meshed region.

    Note that this information is only available for standard solid models
    generated by ACP. The solid stack are computed on request. There is a
    cache to avoid reevaluation of the same stack is retrieved multiple times.
    """

    SOLID_STACK_PROPERTY_FIELD_NAME = "solid_stacks"

    def __init__(self, mesh: MeshedRegion, layup_provider: Operator):
        """
        Get solid stack information from a meshed region.

        Note that this information is only available for standard solid models
        generated by ACP.
        """
        self._mesh = mesh
        self._layup_property_provider = LayupPropertiesProvider(layup_provider, mesh)
        self._mesh_properties_container = (
            layup_provider.outputs.mesh_properties_container.get_data()
        )

        if self._mesh_properties_container:
            self._virtual_thicknesses_field = self._mesh_properties_container.get_field(
                {"MeshPropertyFieldLabel": 2}
            )
        else:
            self._virtual_thicknesses_field = None
        if self.SOLID_STACK_PROPERTY_FIELD_NAME not in self._mesh.available_property_fields:
            raise RuntimeError(
                f"Property field '{self.SOLID_STACK_PROPERTY_FIELD_NAME}' not found in mesh."
                " Please make sure the layup provider has been executed."
            )

        self._solid_stacks_property_field = self._mesh.property_field(
            self.SOLID_STACK_PROPERTY_FIELD_NAME
        )

        # check the number of components. 2 because it is (element_id, level)
        if self._solid_stacks_property_field.ndim != 2:
            raise RuntimeError(
                f"Property field '{self.SOLID_STACK_PROPERTY_FIELD_NAME}' must have 2 components "
                "but it has {self._solid_stacks_property_field.ndim}."
            )

        # Collect all analysis ply fields to process the homogeneous
        # elements (drop-offs and cut-offs) which are linked to an
        # analysis ply but have no layers
        self._analysis_ply_names = get_all_analysis_ply_names(self._mesh)

        # cache to avoid reevaluation
        self._element_id_to_solid_stack_index_map: dict[int, int] = {}
        self._solid_stacks: list[SolidStack] = []

    @property
    def number_of_stacks(self) -> int:
        """Number of solid stacks in the model."""
        return int(self._solid_stacks_property_field.scoping.size)

    def _get_analysis_ply_info_for_homogeneous_element(self, element_id: int) -> dict[str, float]:
        """
        Get the analysis ply info for homogeneous elements.

        Drop-off and cut-off elements origin from layered elements and so the information
        is extracted from the lay-up provider (ACP composite definitions).
        """
        virtual_thickness = None
        if (
            self._virtual_thicknesses_field
            and element_id in self._virtual_thicknesses_field.scoping.ids
        ):
            virtual_thicknesses_array = self._virtual_thicknesses_field.get_entity_data_by_id(
                element_id
            )
            if len(virtual_thicknesses_array) > 0 and virtual_thicknesses_array[0] > 0.0:
                virtual_thickness = virtual_thicknesses_array[0]

        analysis_ply_infos: dict[str, float] = {}
        for ply_name in self._analysis_ply_names:
            analysis_ply_info_provider = AnalysisPlyInfoProvider(self._mesh, ply_name)
            if element_id in analysis_ply_info_provider.ply_element_ids():
                if virtual_thickness:
                    analysis_ply_infos[ply_name] = virtual_thickness
                else:
                    basic_ap_info = analysis_ply_info_provider.basic_info()
                    analysis_ply_infos[ply_name] = basic_ap_info.nominal_thickness

        if virtual_thickness and len(analysis_ply_infos) > 1:
            for k in analysis_ply_infos.keys():
                analysis_ply_infos[k] = virtual_thickness / len(analysis_ply_infos)

        return analysis_ply_infos

    def _build_solid_stack(self, selected_solid_element: int) -> SolidStack | None:
        """
        Create the solid stack by finding the right stack.

        Data of the stack is extracted from the meshed region and
        some additional fields. This fills the cache (self._solid_stacks
        and self._element_id_to_solid_stack_index_map) for later use.

        Returns None if the element is not part of a solid stack.
        """
        for index in range(0, self.number_of_stacks):
            elementary_data = self._solid_stacks_property_field.get_entity_data(index)
            element_ids = [v[0] for v in elementary_data]
            if selected_solid_element in element_ids:
                element_wise_analysis_plies: dict[int, Sequence[str]] = {}
                element_ply_thicknesses: dict[int, Sequence[float]] = {}
                element_wise_levels: dict[int, int] = {}
                for element_id, level in elementary_data:
                    element_wise_levels[int(element_id)] = level
                    ply_ids = self._layup_property_provider.get_analysis_plies(element_id)
                    if ply_ids:
                        element_wise_analysis_plies[element_id] = ply_ids
                        self._element_id_to_solid_stack_index_map[element_id] = len(
                            self._solid_stacks
                        )
                        layer_thicknesses = self._layup_property_provider.get_layer_thicknesses(
                            element_id
                        )
                        if layer_thicknesses is not None:
                            element_ply_thicknesses[int(element_id)] = [
                                float(v) for v in layer_thicknesses
                            ]
                        else:
                            raise RuntimeError("Could not extract the layer thicknesses!")
                    else:
                        ply_infos = self._get_analysis_ply_info_for_homogeneous_element(element_id)
                        if ply_infos:
                            element_wise_analysis_plies[element_id] = list(ply_infos.keys())
                            self._element_id_to_solid_stack_index_map[element_id] = len(
                                self._solid_stacks
                            )
                            element_ply_thicknesses[int(element_id)] = list(ply_infos.values())

                this_stack = SolidStack(
                    element_ids=element_ids,
                    element_wise_analysis_plies=element_wise_analysis_plies,
                    element_ply_thicknesses=element_ply_thicknesses,
                    element_wise_levels=element_wise_levels,
                )
                self._solid_stacks.append(this_stack)
                return this_stack

        return None

    def get_solid_stack(self, element_id: int) -> SolidStack | None:
        """Get the full solid stack for a given element.

        Returns None if the element is not part of a solid stack.
        """
        if element_id in self._element_id_to_solid_stack_index_map:
            solid_stack_index = self._element_id_to_solid_stack_index_map[element_id]
            return self._solid_stacks[solid_stack_index]
        else:
            return self._build_solid_stack(element_id)

    def get_solid_stacks(self, element_ids: Sequence[int]) -> list[SolidStack]:
        """Get unique list of solid stacks for a list of element ids."""
        processed_elements: list[int] = []
        selected_stacks: list[SolidStack] = []
        for e_id in element_ids:
            stack = self.get_solid_stack(e_id)
            if stack:
                if stack.element_ids[0] in processed_elements:
                    # it's enough to check if the first element of the stack is already processed
                    continue
                else:
                    selected_stacks.append(stack)
                    processed_elements.extend(stack.element_ids)
        return selected_stacks
