"""Thermo Elastic 2D Problem."""

from dataclasses import dataclass
from dataclasses import field
from typing import Annotated, Any, ClassVar

from gymnasium import spaces
from matplotlib import colors
import matplotlib.pyplot as plt
import numpy as np
import numpy.typing as npt

from engibench.constraint import bounded
from engibench.constraint import constraint
from engibench.constraint import IMPL
from engibench.constraint import THEORY
from engibench.core import ObjectiveDirection
from engibench.core import OptiStep
from engibench.core import Problem
from engibench.problems.thermoelastic2d.model.fea_model import FeaModel
from engibench.problems.thermoelastic2d.utils import get_res_bounds
from engibench.problems.thermoelastic2d.utils import indices_to_binary_matrix

NELY = NELX = 64
LCI, TRI, RCI, BRI = get_res_bounds(NELX + 1, NELY + 1)
FIXED_ELEMENTS = indices_to_binary_matrix([LCI[21], LCI[32], LCI[43]], NELX + 1, NELY + 1)
FORCE_ELEMENTS_X = indices_to_binary_matrix([BRI[31]], NELX + 1, NELY + 1)
FORCE_ELEMENTS_Y = indices_to_binary_matrix([BRI[31]], NELX + 1, NELY + 1)
HEATSINK_ELEMENTS = indices_to_binary_matrix([LCI[31], LCI[32], LCI[33]], NELX + 1, NELY + 1)


class ThermoElastic2D(Problem[npt.NDArray]):
    r"""Truss 2D integer optimization problem.

    ## Problem Description
    This is 2D topology optimization problem for minimizing weakly coupled thermo-elastic compliance subject to boundary conditions and a volume fraction constraint.

    ## Design space
    The design space is represented by a 2D tensor of continuous design variables in the range [0, 1] that represent the material density at each voxel in the design space.

    ## Objectives
    The objectives are defined and indexed as follows:
    0. `structural_compliance`: Structural compliance to minimize.
    1. `thermal_compliance`: Thermal compliance to minimize.
    2. `volume_fraction`: Volume fraction error to minimize.

    ## Boundary Conditions
    Creating a problem formulation requires defining a python dict with the following info:
    - `fixed_elements`: Encodes a binary NxN matrix of the structurally fixed elements in the domain.
    - `force_elements_x`: Encodes a binary NxN matrix specifying elements that have a structural load in the x-direction.
    - `force_elements_y`: Encodes a binary NxN matrix specifying elements that have a structural load in the y-direction.
    - `heatsink_elements`: Encodes a binary NxN matrix specifying elements that have a heat sink.
    - `volfrac`: Encodes the target volume fraction for the volume fraction constraint.
    - `rmin`: Encodes the filter size used in the optimization routine.
    - `weight`: Allows one to control which objective is optimized for. 1.0 Is pure structural optimization, while 0.0 is pure thermal optimization.

    ## Dataset
    The dataset linked to this problem is on huggingface [Hugging Face Datasets Hub](https://huggingface.co/datasets/IDEALLab/thermoelastic_2d_v0).
    Relevant datapoint fields include:
    - `optimal_design`: An optimized design for the set of boundary conditions
    - `strain`: The strain field for the initial uniform design
    - `vm_stress`: The von Mises stress field for the initial uniform design
    - `fixed_elements`: Encodes a binary NxN matrix of the structurally fixed elements in the domain.
    - `force_elements_x`: Encodes a binary NxN matrix specifying elements that have a structural load in the x-direction.
    - `force_elements_y`: Encodes a binary NxN matrix specifying elements that have a structural load in the y-direction.
    - `heatsink_elements`: Encodes a binary NxN matrix specifying elements that have a heat sink.
    - `volume_fraction`: The volume fraction value of the optimized design
    - `structural_compliance`: The structural compliance of the optimized design
    - `thermal_compliance`: The thermal compliance of the optimized design
    - `nelx`: The number of elements in the x-direction
    - `nely`: The number of elements in the y-direction
    - `volfrac`: The volume fraction target of the optimized design
    - `rmin`: The filter size used in the optimization routine
    - `weight`: The domain weighting used in the optimization routine

    The dataset is generated by enumerating different boundary condition components, and taking the cartesian product of these components. The following components are enumerated:
    - `fixed_elements`: 2 fixed elements that can lie anywhere on the left or top boundary
    - `force_elements_x`: 1 force element that can lie anywhere on the right boundary
    - `force_elements_y`: 1 force element that can lie anywhere on the bottom boundary
    - `heatsink_elements`: N heat sink elements that can lie anywhere on the left or top boundary (N is also enumerated to have different sizes)
    - `volfrac`: The volume fraction target of the optimized design


    ## Simulator
    The evaluation code models the problem as a weakly coupled thermo-elastic problem. The code is written in pure python, and the evaluation is done in a single process.

    ## Lead
    Gabriel Apaza @gapaza
    """

    version = 0
    objectives: tuple[tuple[str, ObjectiveDirection], ...] = (
        ("structural_compliance", ObjectiveDirection.MINIMIZE),
        ("thermal_compliance", ObjectiveDirection.MINIMIZE),
        ("volume_fraction", ObjectiveDirection.MINIMIZE),
    )
    conditions: tuple[tuple[str, Any], ...] = (
        ("fixed_elements", FIXED_ELEMENTS),
        ("force_elements_x", FORCE_ELEMENTS_X),
        ("force_elements_y", FORCE_ELEMENTS_Y),
        ("heatsink_elements", HEATSINK_ELEMENTS),
        ("volfrac", 0.3),
        ("rmin", 1.1),
        ("weight", 0.5),  # 1.0 for pure structural, 0.0 for pure thermal
    )
    design_space = spaces.Box(low=0.0, high=1.0, shape=(NELX, NELY), dtype=np.float32)
    dataset_id = "IDEALLab/thermoelastic_2d_v0"
    container_id = None

    @dataclass
    class Config:
        """Structured representation of configuration parameters for a numerical computation."""

        nelx: ClassVar[Annotated[int, bounded(lower=1).category(THEORY)]] = NELX
        nely: ClassVar[Annotated[int, bounded(lower=1).category(THEORY)]] = NELX
        volfrac: Annotated[float, bounded(lower=0.0, upper=1.0).category(THEORY)] = 0.3
        rmin: Annotated[
            float, bounded(lower=1.0).category(THEORY), bounded(lower=0.0, upper=3.0).warning().category(IMPL)
        ] = 1.1
        weight: Annotated[float, bounded(lower=0.0, upper=1.0).category(THEORY)] = 0.5
        fixed_elements: Annotated[npt.NDArray[np.int64], bounded(lower=0.0, upper=1.0).category(THEORY)] = field(
            default_factory=lambda: FIXED_ELEMENTS
        )
        force_elements_x: Annotated[npt.NDArray[np.int64], bounded(lower=0.0, upper=1.0).category(THEORY)] = field(
            default_factory=lambda: FORCE_ELEMENTS_X
        )
        force_elements_y: Annotated[npt.NDArray[np.int64], bounded(lower=0.0, upper=1.0).category(THEORY)] = field(
            default_factory=lambda: FORCE_ELEMENTS_Y
        )
        heatsink_elements: Annotated[npt.NDArray[np.int64], bounded(lower=0.0, upper=1.0).category(THEORY)] = field(
            default_factory=lambda: HEATSINK_ELEMENTS
        )

        @constraint
        @staticmethod
        def rmin_bound(rmin: float, nelx: int, nely: int) -> None:
            """Constraint for rmin ∈ (0.0, max{ nelx, nely }]."""
            assert 0.0 < rmin <= max(nelx, nely), f"Params.rmin: {rmin} ∉ (0, max(nelx, nely)]"

    def __init__(self) -> None:
        """Initializes the thermoelastic2D problem.

        Args:
            base_directory (str, optional): The base directory for the problem. If None, the current directory is selected.
        """
        super().__init__()
        self.seed = None

    def reset(self, seed: int | None = None) -> None:
        """Resets the simulator and numpy random to a given seed.

        Args:
            seed (int, optional): The seed to reset to. If None, a random seed is used.
        """
        super().reset(seed)

    def simulate(self, design: npt.NDArray, config: dict[str, Any] | None = None) -> npt.NDArray:
        """Simulates the performance of a design topology.

        Args:
            design (np.ndarray): The design to simulate.
            config (dict): A dictionary with configuration (e.g., boundary conditions, filenames) for the simulation.

        Returns:
            dict: The performance of the design - each entry of the dict corresponds to a named objective value.
        """
        boundary_dict = dict(self.conditions)
        for key, value in (config or {}).items():
            if key in boundary_dict:
                boundary_dict[key] = value
        results = FeaModel(plot=False, eval_only=True).run(boundary_dict, x_init=design)
        return np.array([results["structural_compliance"], results["thermal_compliance"], results["volume_fraction"]])

    def optimize(
        self, starting_point: npt.NDArray, config: dict[str, Any] | None = None
    ) -> tuple[np.ndarray, list[OptiStep]]:
        """Optimizes a topology for the current problem. Note that an appropriate starting_point for the optimization is defined by a uniform material distribution equal to the volume fraction constraint.

        Args:
            starting_point (np.ndarray): The starting point for the optimization.
            config (dict): A dictionary with configuration (e.g., boundary conditions, filenames) for the optimization.

        Returns:
            Tuple[np.ndarray, dict]: The optimized design and its performance.
        """
        boundary_dict = dict(self.conditions)
        boundary_dict.update({k: v for k, v in (config or {}).items() if k in dict(self.conditions)})
        results = FeaModel(plot=False, eval_only=False).run(boundary_dict, x_init=starting_point)
        design = np.array(results["design"]).astype(np.float32)
        opti_steps = results["opti_steps"]
        return design, opti_steps

    def render(self, design: np.ndarray, *, open_window: bool = False) -> Any:
        """Renders the design in a human-readable format.

        Args:
            design (np.ndarray): The design to render.
            open_window (bool): If True, opens a window with the rendered design.

        Returns:
            Any: The rendered design.
        """
        fig, ax = plt.subplots(1, 1, figsize=(7, 7))
        ax.imshow(-design, cmap="gray", interpolation="none", norm=colors.Normalize(vmin=-1, vmax=0))
        ax.axis("off")
        plt.tight_layout()
        if open_window is True:
            plt.show()

        return fig

    def random_design(self) -> tuple[npt.NDArray, int]:
        """Samples a valid random design.

        Returns:
            DesignType: The valid random design.
        """
        boundary_dict = dict(self.conditions)
        volfrac = boundary_dict["volfrac"]
        design = volfrac * np.ones((NELX, NELY))
        return design, 0


if __name__ == "__main__":
    # --- Create a new problem
    problem = ThermoElastic2D()
    problem.reset()

    # --- Load the problem dataset
    dataset = problem.dataset
    first_item = dataset["train"][0]
    first_item_design = np.array(first_item["design"])
    problem.render(first_item_design, open_window=True)

    # --- Render the design
    design, _ = problem.random_design()
    problem.render(design, open_window=True)

    # --- Optimize a design ---
    design, _ = problem.random_design()
    design, objectives = problem.optimize(design)
    problem.render(design, open_window=True)

    # --- Evaluate a design ---
    design, _ = problem.random_design()
    print(problem.simulate(design))
