# ruff: noqa: N806, N815, N816
# Disabled variable name conventions

"""Beams 2D problem."""

from copy import deepcopy
from dataclasses import dataclass
from dataclasses import field
from typing import Annotated, Any

from gymnasium import spaces
import numpy as np
import numpy.typing as npt

from engibench.constraint import bounded
from engibench.constraint import constraint
from engibench.constraint import greater_than
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.beams2d.backend import calc_sensitivity
from engibench.problems.beams2d.backend import design_to_image
from engibench.problems.beams2d.backend import image_to_design
from engibench.problems.beams2d.backend import inner_opt
from engibench.problems.beams2d.backend import overhang_filter
from engibench.problems.beams2d.backend import setup
from engibench.problems.beams2d.backend import State


@dataclass
class ExtendedOptiStep(OptiStep):
    """Extended OptiStep to store a single NumPy array representing a density field at a given optimization step."""

    design: npt.NDArray[np.float64] = field(default_factory=lambda: np.array([], dtype=np.float64))


class Beams2D(Problem[npt.NDArray]):
    r"""Beam 2D topology optimization problem.

    ## Problem Description
    Beams2D is a structural topology optimization (TO) problem that optimizes a 2D Messerschmitt-Bölkow-Blohm (MBB) beam under bending. The beam is symmetric about the central vertical axis, with a force applied at the top; only the right half is modeled in our case. Problems are formulated using density-based TO, drawing from an existing Python [implementation](https://github.com/arjendeetman/TopOpt-MMA-Python).

    ## Design Space
    The design space is an array of solid densities in `[0., 1.]` with default image shape `(100, 50)`, where `nelx = 100` and `nely = 50`. Internally, this is represented as a flattened `(5000,)` array. Alternative shapes include `(50, 25)` for faster computation and `(200, 100)` for higher-resolution results. Corresponding datasets for these three resolutions are provided.

    ## Objectives
    The goal is to optimize the distribution of solid material to minimize compliance (equivalently, maximize stiffness) while satisfying constraints on material usage and minimum feature size.

    The objectives are defined and indexed as follows:

    0. `c`: Compliance to minimize.

    ## Conditions
    The following input parameters define the problem conditions:
    - `volfrac`: Desired volume fraction of solid material.
    - `rmin`: Minimum feature length of beam members.
    - `forcedist`: Fractional distance of the downward force from the top-left (default) to the top-right corner.
    - `overhang_constraint`: Boolean flag to enable a 45-degree overhang constraint for manufacturability.

    ## Simulator
    Compliance `c` is calculated using:
    ```python
    c = ((Emin + xPrint**penal * (Emax - Emin)) * ce).sum()
    ```

    where `xPrint` is the current true density field, `penal` is the penalization factor (e.g., 3.0), and `ce` is the element-wise strain energy density.

    ## Dataset
    Three datasets, each corresponding to a different resolution, are available on the [Hugging Face Datasets Hub](https://huggingface.co/datasets/IDEALLab/beams_2d).

    ### v0

    #### Fields
    Each dataset contains:
    - Optimized beam structures,
    - The corresponding condition parameters,
    - Objective values (compliance),
    - Full optimization histories (for advanced use).

    #### Creation Method
    Datasets were generated by uniformly sampling the condition space. The resolutions used are:
    - `(50, 25)`
    - `(100, 50)`
    - `(200, 100)`

    A more comprehensive description of the creation method can be found in the [README](https://github.com/IDEALLab/EngiBench/tree/main/engibench/problems/beams2d).

    ## References
    If you use this problem in your research, please cite the following paper:
    E. Andreassen, A. Clausen, M. Schevenels, B. S. Lazarov, and O. Sigmund, “Efficient topology optimization in MATLAB using 88 lines of code,” in Structural and Multidisciplinary Optimization, vol. 43, pp. 1-16, 2011.

    ## Lead
    Arthur Drake @arthurdrake1
    """

    version = 0
    objectives: tuple[tuple[str, ObjectiveDirection]] = (("c", ObjectiveDirection.MINIMIZE),)
    nelx = 100
    nely = 50
    conditions: tuple[tuple[str, Any], ...] = (
        ("volfrac", 0.35),
        ("rmin", 2.0),
        ("forcedist", 0.0),
        ("overhang_constraint", False),
    )
    design_space = spaces.Box(low=0.0, high=1.0, shape=(nely, nelx), dtype=np.float64)
    dataset_id = f"IDEALLab/beams_2d_{nely}_{nelx}_v{version}"
    container_id = None

    def __init__(self, config: dict[str, Any] | None = None):
        """Initializes the Beams2D problem.

        Args:
            config (dict): A dictionary with configuration (e.g., boundary conditions) for the simulation.
        """
        super().__init__()

        config = config or {}

        # Replace the conditions with any new configs passed in
        self.conditions = tuple((key, config.get(key, value)) for key, value in self.conditions)
        self.__st = State()
        if "nely" in config and "nelx" in config:
            self.nelx = config["nelx"]
            self.nely = config["nely"]
            self.design_space = spaces.Box(low=0.0, high=1.0, shape=(self.nely, self.nelx), dtype=np.float64)
            self.dataset_id = f"IDEALLab/beams_2d_{self.nely}_{self.nelx}_v{self.version}"

        # Default values of Config depend on nelx, nely => define it in __init__
        @dataclass
        class Config:
            """Structured representation of configuration parameters for a numerical computation."""

            nelx: Annotated[
                int, bounded(lower=1).category(THEORY), bounded(lower=10, upper=1000).warning().category(IMPL)
            ] = self.nelx
            nely: Annotated[
                int, bounded(lower=1).category(THEORY), bounded(lower=10, upper=1000).warning().category(IMPL)
            ] = self.nely
            penal: Annotated[
                float, bounded(lower=1.0).category(IMPL), bounded(lower=2.0, upper=5.0).category(IMPL).warning()
            ] = 3.0
            max_iter: Annotated[
                int, bounded(lower=0).category(THEORY), bounded(lower=1, upper=1000).category(IMPL).warning()
            ] = 100
            volfrac: Annotated[
                float,
                bounded(lower=0.0, upper=1.0).category(THEORY),
                bounded(lower=0.1, upper=0.9).warning().category(IMPL),
            ] = 0.35
            rmin: Annotated[
                float, greater_than(0.0).category(THEORY), bounded(lower=1.0, upper=10.0).category(IMPL).warning()
            ] = 2.0
            forcedist: Annotated[float, bounded(lower=0.0, upper=1.0).category(THEORY)] = 0.0
            overhang_constraint: bool = False

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

        self.Config = Config

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

        Args:
            design (np.ndarray): The design to simulate.
            ce: (np.ndarray, optional): If applicable, the pre-calculated sensitivity of the current design.
            config (dict): A dictionary with configuration (e.g., boundary conditions) for the simulation.

        Returns:
            npt.NDArray: The performance of the design in terms of compliance.
        """
        # This condition is needed to convert user-provided designs (images) to flat arrays. Normally does not apply, i.e., during optimization.
        if len(design.shape) > 1:
            design = image_to_design(design)

        base_config = {"nelx": self.nelx, "nely": self.nely, "penal": 3.0, **dict(self.conditions), **(config or {})}

        # Assumes ndof is initialized as 0. This is a check to see if setup has run yet.
        # If setup has run, skips the process for repeated simulations during optimization.
        if self.__st.ndof == 0:
            self.__st = setup(base_config)

        if ce is None:
            ce = calc_sensitivity(design, st=self.__st, cfg=base_config)
        c = (
            (self.__st.Emin + design ** base_config["penal"] * (self.__st.Emax - self.__st.Emin)) * ce
        ).sum()  # compliance (objective)
        return np.array([c])

    def optimize(
        self, starting_point: npt.NDArray | None = None, config: dict[str, Any] | None = None
    ) -> tuple[np.ndarray, list[ExtendedOptiStep]]:
        """Optimizes the design of a beam.

        Args:
            starting_point (npt.NDArray or None): The design to begin warm-start optimization from (optional).
            config (dict): A dictionary with configuration (e.g., boundary conditions) for the optimization.

        Returns:
            Tuple[np.ndarray, dict]: The optimized design and its performance.
        """
        base_config = {
            "nelx": self.nelx,
            "nely": self.nely,
            "max_iter": 100,
            "penal": 3.0,
            **dict(self.conditions),
            **(config or {}),
        }

        self.__st = setup(base_config)

        # Returns the full history of the optimization instead of just the last step
        optisteps_history = []

        if starting_point is None:
            xPhys = x = base_config["volfrac"] * np.ones(base_config["nely"] * base_config["nelx"], dtype=float)
            dc = np.zeros(base_config["nely"] * base_config["nelx"])
            dv = np.zeros(base_config["nely"] * base_config["nelx"])
        else:
            starting_point = image_to_design(starting_point)
            assert starting_point is not None
            xPhys = x = deepcopy(starting_point)
            ce = calc_sensitivity(starting_point, st=self.__st, cfg=base_config)
            dc = (
                -base_config["penal"] * starting_point ** (base_config["penal"] - 1) * (self.__st.Emax - self.__st.Emin)
            ) * ce
            dv = np.ones(base_config["nely"] * base_config["nelx"])

        xPrint, _, _ = overhang_filter(xPhys, base_config, dc, dv)
        loop, change = (0, 1.0)

        while change > self.__st.min_change and loop < base_config["max_iter"]:
            ce = calc_sensitivity(xPrint, st=self.__st, cfg=base_config)
            c = self.simulate(xPrint, ce=ce, config=base_config)

            # Record the current state in optisteps_history
            current_step = ExtendedOptiStep(obj_values=np.array(c), step=loop)
            current_step.design = np.array(xPrint)
            optisteps_history.append(current_step)

            loop += 1

            dc = (-base_config["penal"] * xPrint ** (base_config["penal"] - 1) * (self.__st.Emax - self.__st.Emin)) * ce
            dv = np.ones(base_config["nely"] * base_config["nelx"])
            xPrint, dc, dv = overhang_filter(xPhys, base_config, dc, dv)  # MATLAB implementation

            dc = np.asarray(self.__st.H * (dc[np.newaxis].T / self.__st.Hs))[:, 0]
            dv = np.asarray(self.__st.H * (dv[np.newaxis].T / self.__st.Hs))[:, 0]
            # Ensure dc remains nonpositive
            dc = np.clip(dc, None, 0.0)

            xnew, xPhys, xPrint = inner_opt(x, self.__st, dc, dv, base_config)
            # Compute the change by the inf. norm
            change = np.linalg.norm(  # type: ignore[assignment]
                xnew.reshape(base_config["nelx"] * base_config["nely"], 1)
                - x.reshape(base_config["nelx"] * base_config["nely"], 1),
                np.inf,
            )
            x = deepcopy(xnew)

        return design_to_image(xPrint, base_config["nelx"], base_config["nely"]), optisteps_history

    def reset(self, seed: int | None = None, **kwargs) -> None:
        r"""Reset numpy random to a given seed.

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

    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.
        """
        import matplotlib.pyplot as plt
        import seaborn as sns

        fig, ax = plt.subplots(figsize=(8, 4))
        sns.heatmap(design, cmap="coolwarm", ax=ax, vmin=0, vmax=1)

        if open_window:
            plt.show()
        return fig, ax

    def random_design(self, dataset_split: str = "train", design_key: str = "optimal_design") -> tuple[npt.NDArray, int]:
        """Samples a valid random design.

        Args:
            dataset_split (str): The key for the dataset to sample from.
            design_key (str): The key for the design to sample from.

        Returns:
            Tuple of:
                np.ndarray: The valid random design.
                int: The random index selected.
        """
        rnd = self.np_random.integers(low=0, high=len(self.dataset[dataset_split]), dtype=int)
        return np.array(self.dataset[dataset_split][design_key][rnd]), rnd


if __name__ == "__main__":
    # Provides a way to instantiate the problem without having to pass configs to optimize or simulate later.
    # Possible sets of nely and nelx: (25, 50), (50, 100), and (100, 200)
    # If a new nely and nelx are not passed in, uses the default conditions.

    problem = Beams2D()
    problem.reset(seed=0)

    print(f"Loading dataset for nely={problem.nely}, nelx={problem.nelx}.")
    dataset = problem.dataset

    # Example of getting the training set
    optimal_train = dataset["train"]["optimal_design"]
    c_train = dataset["train"]["c"]
    params_train = dataset["train"].select_columns(problem.conditions_keys)

    # Get design and conditions from the dataset, render design
    # Note that here, we override any previous configs to re-optimize the same design as a test case.
    design, idx = problem.random_design()
    config = params_train[idx]
    compliance = c_train[idx]
    fig, ax = problem.render(design, open_window=True)

    print(f"Verifying compliance via simulation. Reference value: {compliance:.4f}")

    try:
        c_ref = problem.simulate(design, config=config)[0]
        print(f"Calculated compliance: {c_ref:.4f}")
    except ArithmeticError:
        print("Failed to calculate compliance for upscaled design.")

    # Sample Optimization
    print("\nNow conducting a sample optimization with the given configs:", config)

    # NOTE: optimal_design and optisteps_history[-1].stored_design are interchangeable.
    optimal_design, optisteps_history = problem.optimize(config=config)
    print(f"Final compliance: {optisteps_history[-1].obj_values[0]:.4f}")
    print(f"Final design volume fraction: {optimal_design.sum() / (np.prod(optimal_design.shape)):.4f}")

    fig, ax = problem.render(optimal_design, open_window=True)
