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

"""Beams 2D problem.

This code has been adapted from the Python implementation by Niels Aage and Villads Egede Johansen: https://github.com/arjendeetman/TopOpt-MMA-Python
"""

import dataclasses
from typing import Any

import cvxopt
import cvxopt.cholmod
import numpy as np
import numpy.typing as npt
from scipy.sparse import coo_matrix
from scipy.sparse import csc_array
from scipy.sparse import sparray


@dataclasses.dataclass
class State:
    """A structured representation of scalars and matrices for optimization and simulation.

    Attributes:
        Emin: (float) Minimum possible stiffness (1e-9 by default).
        Emax: (float) Maximum possible stiffness (1 by default).
        min_change (float): Minimum change in terms of design variables between two consecutive designs to continue optimization (0.025 by default).
        min_ratio (float): Parameter determining when the bisection search on the Lagrange multiplier should stop (1e-3 by default).
        ndof (int): Number of degrees of freedom.
        edofMat (np.ndarray): Element degrees of freedom mapping.
        iK (np.ndarray): Row indices for stiffness matrix.
        jK (np.ndarray]): Column indices for stiffness matrix.
        H (csc_array): Filter matrix.
        Hs (np.ndarray): Filter normalization factor.
        dofs (np.ndarray): Degrees of freedom indices.
        fixed (np.ndarray): Fixed degrees of freedom.
        free (np.ndarray): Free degrees of freedom.
        f (np.ndarray): Force vector.
        u (np.ndarray): Displacement vector.
        KE (np.ndarray): Stiffness matrix.
    """

    # Non-editable Constants
    Emin: float = 1e-9
    Emax: float = 1.0
    min_change: float = 0.025
    min_ratio: float = 1.0e-3

    # Items calculated for optimization and simulation
    ndof: int = 0
    edofMat: np.ndarray = dataclasses.field(default_factory=lambda: np.array([]))
    iK: np.ndarray = dataclasses.field(default_factory=lambda: np.array([]))
    jK: np.ndarray = dataclasses.field(default_factory=lambda: np.array([]))
    H: csc_array = dataclasses.field(default_factory=lambda: csc_array((0, 0)))
    Hs: np.ndarray = dataclasses.field(default_factory=lambda: np.array([]))
    dofs: np.ndarray = dataclasses.field(default_factory=lambda: np.array([]))
    fixed: np.ndarray = dataclasses.field(default_factory=lambda: np.array([]))
    free: np.ndarray = dataclasses.field(default_factory=lambda: np.array([]))
    f: np.ndarray = dataclasses.field(default_factory=lambda: np.array([]))
    u: np.ndarray = dataclasses.field(default_factory=lambda: np.array([]))
    KE: np.ndarray = dataclasses.field(default_factory=lambda: np.array(lk()))

    def get(self, keys: list[str]) -> dict[str, Any]:
        """Retrieve a subset of state values based on a list of keys.

        Args:
            keys (List[str]): List of state names to retrieve.

        Returns:
            Dict[str, Any]: Dictionary of requested state names and values.
        """
        return {key: getattr(self, key) for key in keys if hasattr(self, key)}

    def to_dict(self) -> dict[str, Any]:
        """Convert the dataclass to a dictionary representation.

        Returns:
            Dict[str, Any]: Dictionary containing all state values.
        """
        return dataclasses.asdict(self)

    @classmethod
    def new(cls, nelx: int, nely: int, rmin: float, forcedist: float) -> "State":
        r"""Set up the scalars and matrices for optimization or simulation.

        Args:
            nelx: Number of elements in x direction.
            nely: Number of elements in y direction.
            rmin: ...
            forcedist: ...

        Returns:
            State object with the relevant scalars and matrices used in optimization and simulation.
        """
        edofMat = edof_mat(nelx, nely)

        H = h_mat(nelx, nely, rmin)

        # BC's and support
        ndof = 2 * (nelx + 1) * (nely + 1)
        dofs = np.arange(ndof)
        fixed = np.union1d(dofs[0 : 2 * (nely + 1) : 2], np.array([2 * (nelx + 1) * (nely + 1) - 1]))

        # Solution and RHS vectors
        f = np.zeros((ndof, 1))
        # Set load at the specified fractional distance (p.forcedist) from the top-left (default) to the top-right corner.
        f[int(1 + (2 * forcedist * nelx) * (nely + 1)), 0] = -1

        return cls(
            ndof=ndof,
            edofMat=edofMat,
            iK=np.kron(edofMat, np.ones((8, 1))).flatten(),
            jK=np.kron(edofMat, np.ones((1, 8))).flatten(),
            H=H,
            Hs=H.sum(1),
            dofs=dofs,
            fixed=fixed,
            free=np.setdiff1d(dofs, fixed),
            f=f,
            u=np.zeros((ndof, 1)),
        )


def image_to_design(im: npt.NDArray) -> npt.NDArray:
    r"""Flatten the 2D image(s) to 1D vector(s).

    Args:
        im (npt.NDArray): The image(s) to convert.

    Returns:
        npt.NDArray: The transformed vector(s).
    """
    return np.swapaxes(im, -2, -1).reshape(*im.shape[:-2], -1)


def design_to_image(x: npt.NDArray, nelx: int = 100, nely: int = 50) -> npt.NDArray:
    r"""Reshape the 1D vector(s) into 2D image(s).

    Args:
        x (npt.NDArray): The design(s) to convert.
        nelx (int): Width of the problem domain.
        nely (int): Height of the problem domain.

    Returns:
        npt.NDArray: The transformed image(s).
    """
    return np.swapaxes(x.reshape(*x.shape[:-1], nelx, nely), -2, -1)


def lk() -> npt.NDArray:
    r"""Set up the stiffness matrix.

    Returns:
        KE (npt.NDArray): The stiffness matrix.
    """
    E = 1  # 1
    nu = 0.3  # 0.3
    k = np.array(
        [
            1 / 2 - nu / 6,
            1 / 8 + nu / 8,
            -1 / 4 - nu / 12,
            -1 / 8 + 3 * nu / 8,
            -1 / 4 + nu / 12,
            -1 / 8 - nu / 8,
            nu / 6,
            1 / 8 - 3 * nu / 8,
        ]
    )
    return (
        E
        / (1 - nu**2)
        * np.array(
            [
                [k[0], k[1], k[2], k[3], k[4], k[5], k[6], k[7]],
                [k[1], k[0], k[7], k[6], k[5], k[4], k[3], k[2]],
                [k[2], k[7], k[0], k[5], k[6], k[3], k[4], k[1]],
                [k[3], k[6], k[5], k[0], k[7], k[2], k[1], k[4]],
                [k[4], k[5], k[6], k[7], k[0], k[1], k[2], k[3]],
                [k[5], k[4], k[3], k[2], k[1], k[0], k[7], k[6]],
                [k[6], k[3], k[4], k[1], k[2], k[7], k[0], k[5]],
                [k[7], k[2], k[1], k[4], k[3], k[6], k[5], k[0]],
            ]
        )
    )


def calc_sensitivity(design: npt.NDArray, st: State, cfg: dict[str, Any] | None = None) -> npt.NDArray:
    """Simulates the performance of a beam design. Assumes the State object is already set up.

    Args:
        design (np.ndarray): The design to simulate.
        st: State object with needed vectors/matrices for the simulation.
        cfg (dict): A dictionary with configurations (e.g., boundary conditions) for the simulation.

    Returns:
        npt.NDArray: The sensitivity of the current design.
    """
    cfg = cfg or {}
    sK = ((st.KE.flatten()[np.newaxis]).T * (st.Emin + design ** cfg["penal"] * (st.Emax - st.Emin))).flatten(order="F")
    K: csc_array = coo_matrix((sK, (st.iK, st.jK)), shape=(st.ndof, st.ndof)).tocsc()
    assert K.shape is not None
    m = K.shape[0]
    keep = np.delete(np.arange(0, m), st.fixed)
    K = K[keep, :]
    keep = np.delete(np.arange(0, m), st.fixed)
    K = K[:, keep].tocoo()
    # Solve system
    K = cvxopt.spmatrix(K.data, K.row.astype(int), K.col.astype(int))
    B = cvxopt.matrix(st.f[st.free, 0])
    cvxopt.cholmod.linsolve(K, B)
    st.u[st.free, 0] = np.array(B)[:, 0]

    ############################################################################################################
    # Sensitivity
    ce = (
        np.dot(st.u[st.edofMat].reshape(cfg["nelx"] * cfg["nely"], 8), st.KE)
        * st.u[st.edofMat].reshape(cfg["nelx"] * cfg["nely"], 8)
    ).sum(1)
    return np.array(ce)


def edof_mat(nelx: int, nely: int) -> npt.NDArray[np.float64]:
    """Assemble the element degrees of freedom mapping matrix used in :class:`State`."""
    n1 = ((nely + 1) * np.arange(nelx)[:, None] + np.arange(nely)).ravel()
    n2 = n1 + nely + 1
    return np.array(
        [
            2 * n1 + 2,
            2 * n1 + 3,
            2 * n2 + 2,
            2 * n2 + 3,
            2 * n2,
            2 * n2 + 1,
            2 * n1,
            2 * n1 + 1,
        ]
    ).T


def h_mat(nelx: int, nely: int, rmin: float) -> sparray:
    """Assemble the filter matrix used in :class:`State`."""
    nfilter = int(nelx * nely * ((2 * (np.ceil(rmin) - 1) + 1) ** 2))
    iH = np.zeros(nfilter)
    jH = np.zeros(nfilter)
    sH = np.zeros(nfilter)
    cc = 0

    r = np.ceil(rmin).astype(np.int64)
    for i in range(nelx):
        for j in range(nely):
            row = i * nely + j
            kk1 = np.maximum(i - (r - 1), 0)
            kk2 = np.minimum(i + r, nelx)
            ll1 = np.maximum(j - (r - 1), 0)
            ll2 = np.minimum(j + r, nely)
            for k in range(kk1, kk2):
                for l in range(ll1, ll2):
                    col = k * nely + l
                    fac = rmin - np.hypot(i - k, j - l)
                    iH[cc] = row
                    jH[cc] = col
                    sH[cc] = np.maximum(0.0, fac)
                    cc += 1
    # Finalize assembly and convert to csc format
    return coo_matrix((sH, (iH, jH)), shape=(nelx * nely, nelx * nely)).tocsc()


def inner_opt(
    x: npt.NDArray,
    st: State,
    dc: npt.NDArray,
    dv: npt.NDArray,
    cfg: dict[str, Any] | None = None,
) -> tuple[npt.NDArray, npt.NDArray, npt.NDArray]:
    """Inner optimization loop: Lagrange Multiplier Optimization.

    Args:
        x: (npt.NDArray) The current density field during optimization.
        st: State object with needed vectors/matrices for the optimization.
        dc: (npt.NDArray) The sensitivity field wrt. compliance.
        dv: (npt.NDArray) The sensitivity field wrt. volume fraction.
        cfg (dict): A dictionary with configurations (e.g., boundary conditions) for the optimization.

    Returns:
        Tuple of:
            npt.NDArray: The raw density field
            npt.NDArray: The processed density field (without overhang constraint)
            npt.NDArray: The processed density field (with overhang constraint if applicable)
    """
    cfg = cfg or {}
    nelx, nely = cfg["nelx"], cfg["nely"]
    overhang_constraint = cfg["overhang_constraint"]
    # Optimality criteria
    l1, l2, move = (0.0, 1e9, 0.2)
    # reshape to perform vector operations
    xnew = np.zeros(nelx * nely)
    xPhys = np.zeros((nelx, nely))
    xPrint = np.zeros(nelx * nely)

    while l1 + l2 > 0 and (l2 - l1) / (l1 + l2) > st.min_ratio:
        lmid = 0.5 * (l2 + l1)
        if lmid > 0:
            xnew = np.maximum(
                0.0, np.maximum(x - move, np.minimum(1.0, np.minimum(x + move, x * np.sqrt(-dc / dv / lmid))))
            )
        else:
            xnew = np.maximum(0.0, np.maximum(x - move, np.minimum(1.0, x + move)))

        # Filter design variables
        xPhys = np.asarray(st.H * xnew[np.newaxis].T / st.Hs)[:, 0].reshape((nelx, nely))
        xPrint = overhang_filter_x(xPhys) if overhang_constraint else xPhys.ravel()

        if xPrint.sum() > cfg["volfrac"] * nelx * nely:
            l1 = lmid
        else:
            l2 = lmid

        # Ensures this loop does not become stuck due to abs(l2 - l1) converging to near 0
        if abs(l2 - l1) < np.finfo(float).eps:
            break

    return (xnew, xPhys, xPrint)


P = 40
SHIFT = 100 * (np.finfo(float).tiny) ** (1 / P)
Ns = 3
xi_0 = 0.5
Q = P + np.log(Ns) / np.log(xi_0)


def compute_xi(
    x: npt.NDArray[np.float64],
) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64], npt.NDArray[np.float64], npt.NDArray[np.float64]]:
    """Shared part between :function:`overhang_filter_x` and  :function:`overhang_filter_d`."""
    ep = 1e-4
    nelx, nely = x.shape

    x = design_to_image(x.ravel(), nelx, nely)

    BACKSHIFT = 0.95 * (Ns ** (1 / Q)) * (SHIFT ** (P / Q))
    xi = np.zeros(x.shape)
    Xi = np.zeros(x.shape)
    keep = np.zeros(x.shape)
    sq = np.zeros(x.shape)

    xi[nely - 1, :] = x[nely - 1, :].copy()
    for i in reversed(range(nely - 1)):
        cbr = np.array([0, *list(xi[i + 1, :]), 0]) + SHIFT
        keep[i, :] = cbr[:nelx] ** P + cbr[1 : nelx + 1] ** P + cbr[2:] ** P
        Xi[i, :] = keep[i, :] ** (1 / Q) - BACKSHIFT
        sq[i, :] = np.sqrt((x[i, :] - Xi[i, :]) ** 2 + ep)
        xi[i, :] = 0.5 * ((x[i, :] + Xi[i, :]) - sq[i, :] + np.sqrt(ep))
    return xi, Xi, sq, keep


def overhang_filter_x(x: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
    """Topology Optimization (TO) filter.

    Args:
        x: (npt.NDArray) The current density field during optimization.

    Returns:
        npt.NDArray: The updated design.
    """
    xi, *_ = compute_xi(x)
    return image_to_design(xi)


def overhang_filter_d(
    x: npt.NDArray[np.float64],
    dc: npt.NDArray[np.float64],
    dv: npt.NDArray[np.float64],
) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64], npt.NDArray[np.float64]]:
    """Topology Optimization (TO) filter.

    Args:
        x: (npt.NDArray) The current density field during optimization.
        dc: (npt.NDArray) The sensitivity field wrt. compliance.
        dv: (npt.NDArray) The sensitivity field wrt. volume fraction.

    Returns:
        Tuple[npt.NDArray, npt.NDArray, npt.NDArray]: The updated design, sensitivity dc, and sensitivity dv, respectively.
    """
    xi, Xi, sq, keep = compute_xi(x)
    nelx, nely = x.shape
    x = design_to_image(x.ravel(), nelx, nely)
    nSens = 2  # dc and dv (hard-coded)

    dc_copy = design_to_image(dc, nelx, nely)
    dv_copy = design_to_image(dv, nelx, nely)
    dfxi = [np.array(dc_copy), np.array(dv_copy)]
    dfx = [np.array(dc_copy), np.array(dv_copy)]
    lamb = np.zeros((nSens, nelx))
    for i in range(nely - 1):
        dsmindx = 0.5 * (1 - (x[i, :] - Xi[i, :]) / sq[i, :])
        dsmindXi = 1 - dsmindx
        cbr = np.array([0, *list(xi[i + 1, :]), 0]) + SHIFT

        dmx = np.zeros((Ns, nelx))
        for j in range(Ns):
            dmx[j, :] = (P / Q) * (keep[i, :] ** (1 / Q - 1)) * (cbr[j : nelx + j] ** (P - 1))

        qi = np.ravel([[i] * 3 for i in range(nelx)])
        qj = qi + [-1, 0, 1] * nelx
        qs = np.ravel(dmx.T)

        dsmaxdxi = coo_matrix((qs[1:-1], (qi[1:-1], qj[1:-1]))).tocsc()
        for k in range(nSens):
            dfx[k][i, :] = dsmindx * (dfxi[k][i, :] + lamb[k, :])
            lamb[k, :] = ((dfxi[k][i, :] + lamb[k, :]) * dsmindXi) @ dsmaxdxi

    i = nely - 1
    for k in range(nSens):
        dfx[k][i, :] = dfx[k][i, :] + lamb[k, :]

    dc, dv = dfx
    dc, dv = image_to_design(dc), image_to_design(dv)

    return (image_to_design(xi), dc, dv)
