"""The plate as evaluated on the grid: material, thickness field, stress-free
shape, and the derived quantities both models need (masses, the strike patch,
the listener terms). A ``Plate`` is usually produced by ``Design.evaluate``;
``Plate.uniform`` builds one directly for tests and experiments."""

from __future__ import annotations

import math
from dataclasses import dataclass, replace
from functools import cached_property

import numpy as np

from .constants import RHO_AIR, ROOM_TEMPERATURE, SPEED_OF_SOUND
from .grid import Grid
from .material import Material
from .params import Listener, Strike


@dataclass(frozen=True)
class Patch:
    """The plate nodes that a strike acts on."""

    mask: np.ndarray
    """``(ny, nx)`` bool: the nodes of the patch."""
    nearest: tuple[int, int]
    """``(j, i)`` of the plate node nearest the strike point."""
    mass: float
    """Total lumped mass of the patch [kg]."""
    cx: float
    """Mass-weighted center of the patch [m]."""
    cy: float


@dataclass(frozen=True)
class Plate:
    grid: Grid
    material: Material
    h: np.ndarray
    """Thickness at the nodes [m], ``(ny, nx)``, zero outside the plate."""
    w0: np.ndarray
    """Stress-free height at the nodes [m], mean plane removed, zero outside the plate."""
    Lx: float
    """Bounding box [m]."""
    Ly: float
    temperature: float = ROOM_TEMPERATURE
    """Temperature [K], for the thermoelastic damping."""
    name: str = "plate"

    def __post_init__(self) -> None:
        if self.h.shape != self.grid.shape or self.w0.shape != self.grid.shape:
            raise ValueError("h and w0 must have the grid's shape (ny, nx)")
        if (self.h[self.grid.mask] <= 0).any():
            raise ValueError("the thickness must be positive on the plate")

    @classmethod
    def uniform(cls, grid: Grid, material: Material, thickness: float, profile: np.ndarray | None = None, name: str = "uniform") -> "Plate":
        """A plate of uniform thickness on a grid, optionally with a stress-free profile (its mean plane is removed)."""
        h = np.where(grid.mask, thickness, 0.0)
        w0 = np.zeros(grid.shape) if profile is None else np.where(grid.mask, remove_plane(np.asarray(profile, float), grid), 0.0)
        return cls(grid, material, h, w0, (grid.nx - 1) * grid.hx, (grid.ny - 1) * grid.hy, name=name)

    def with_thickness(self, h: np.ndarray) -> "Plate":
        """The same plate with another thickness field (masked to the plate)."""
        return replace(self, h=np.where(self.grid.mask, h, 0.0))

    # ---- masses and geometry ----

    @cached_property
    def node_mass(self) -> np.ndarray:
        """Lumped mass of each node, rho h dA times its weight [kg]."""
        return self.material.rho * self.h * self.grid.node_area

    @cached_property
    def mass(self) -> float:
        """Total mass [kg]."""
        return float(self.node_mass.sum())

    @cached_property
    def center(self) -> tuple[float, float]:
        """Mass-weighted center [m]."""
        m = self.node_mass
        return float((m * self.grid.X).sum() / self.mass), float((m * self.grid.Y).sum() / self.mass)

    @cached_property
    def inertia(self) -> tuple[float, float]:
        """Moments of inertia ``(Ix, Iy)`` about the in-plane axes through the center [kg m^2]."""
        cx, cy = self.center
        m = self.node_mass
        return float((m * (self.grid.Y - cy) ** 2).sum()), float((m * (self.grid.X - cx) ** 2).sum())

    @property
    def area(self) -> float:
        return self.grid.area

    @cached_property
    def mean_thickness(self) -> float:
        """Area-weighted mean thickness [m]."""
        w = self.grid.node_weight
        return float((w * self.h).sum() / w.sum())

    @cached_property
    def thickness_range(self) -> tuple[float, float]:
        hh = self.h[self.grid.mask]
        return float(hh.min()), float(hh.max())

    @cached_property
    def cell_thickness(self) -> np.ndarray:
        """Thickness of each cell, the mean of its four corners, ``(ny - 1, nx - 1)``."""
        h = self.h
        return 0.25 * (h[:-1, :-1] + h[:-1, 1:] + h[1:, 1:] + h[1:, :-1])

    @cached_property
    def is_flat(self) -> bool:
        """Whether the profile vanishes, in which case bending and in-plane motion decouple in the linear model."""
        return not np.any(self.w0)

    @cached_property
    def profile_stats(self) -> tuple[float, float]:
        """Peak-to-peak height and largest slope of the profile over the plate nodes."""
        g, w0, m = self.grid, self.w0, self.grid.mask
        left, right, down, up = g.neighbor_masks()
        sx = np.zeros(g.shape)
        sy = np.zeros(g.shape)
        wp = np.pad(w0, 1)
        both = left & right
        sx[both] = ((wp[1:-1, 2:] - wp[1:-1, :-2]) / (2 * g.hx))[both]
        only = right & ~left
        sx[only] = ((wp[1:-1, 2:] - w0) / g.hx)[only]
        only = left & ~right
        sx[only] = ((w0 - wp[1:-1, :-2]) / g.hx)[only]
        both = down & up
        sy[both] = ((wp[2:, 1:-1] - wp[:-2, 1:-1]) / (2 * g.hy))[both]
        only = up & ~down
        sy[only] = ((wp[2:, 1:-1] - w0) / g.hy)[only]
        only = down & ~up
        sy[only] = ((w0 - wp[:-2, 1:-1]) / g.hy)[only]
        v = w0[m]
        return float(v.max() - v.min()), float(np.hypot(sx, sy)[m].max())

    def dispersion_frequency(self, fraction: float, h: float | None = None) -> float:
        """Frequency [Hz] at which grid dispersion lowers a bending partial by
        ``fraction`` (a plane wave along a grid axis has the discrete frequency
        ratio (2 - 2 cos kh) / (kh)^2), at the thickness ``h`` (the mean by default)."""
        lo, hi = 0.0, math.pi
        for _ in range(60):
            mid = (lo + hi) / 2
            ratio = (2 - 2 * math.cos(mid)) / (mid * mid)
            if 1 - ratio < fraction:
                lo = mid
            else:
                hi = mid
        k = lo / max(self.grid.hx, self.grid.hy)
        return self.material.bending_constant(self.mean_thickness if h is None else h) * k * k / (2 * math.pi)

    # ---- strike and listener ----

    def strike_point(self, strike: Strike) -> tuple[float, float]:
        """The strike point in plate coordinates [m]."""
        return (strike.x - 0.5) * self.Lx, (strike.y - 0.5) * self.Ly

    def patch_around(self, x: float, y: float, radius: float) -> Patch:
        """The plate nodes within ``radius`` of ``(x, y)``, always including the nearest one."""
        g = self.grid
        d = np.hypot(g.X - x, g.Y - y)
        d = np.where(self.node_mass > 0, d, np.inf)
        nearest = np.unravel_index(int(np.argmin(d)), g.shape)
        mask = d <= radius
        mask[nearest] = True
        m = self.node_mass[mask]
        mass = float(m.sum())
        return Patch(mask, (int(nearest[0]), int(nearest[1])), mass, float((m * g.X[mask]).sum() / mass), float((m * g.Y[mask]).sum() / mass))

    def strike_patch(self, strike: Strike) -> Patch:
        x, y = self.strike_point(strike)
        return self.patch_around(x, y, strike.radius)

    def listener_terms(self, listener: Listener) -> tuple[np.ndarray, np.ndarray]:
        """Per-node terms of the Rayleigh integral: the weight rho_air dA / (2 pi r)
        and the delay r / c [s] to the listener, ``(ny, nx)`` each, zero outside
        the plate. The profile is included in the node positions."""
        g = self.grid
        lx, ly, lz = listener.position
        r = np.sqrt((lx - g.X) ** 2 + (ly - g.Y) ** 2 + (lz - self.w0) ** 2)
        area = g.node_area
        weight = np.where(area > 0, RHO_AIR * area / (2 * math.pi * np.maximum(r, 1e-3)), 0.0)
        delay = np.where(area > 0, r / SPEED_OF_SOUND, 0.0)
        return weight, delay

    # ---- summary ----

    def summary(self) -> str:
        """A few lines describing the plate: box, material, grid, thickness, mass, profile, loss constants."""
        from .loss import LN1000, coincidence_frequency, estimated_rate, plate_geometry

        g, m = self.grid, self.material
        hmin, hmax = self.thickness_range
        lines = [
            f"{self.name}: {self.Lx * 1e3:.0f} x {self.Ly * 1e3:.0f} mm box, area {self.area * 1e4:.0f} cm^2, mass {self.mass * 1e3:.0f} g",
            f"material: E = {m.E / 1e9:.1f} GPa, nu = {m.nu}, rho = {m.rho:.0f} kg/m^3, loss factor {m.loss_factor:g}"
            + (", thermal constants given" if m.thermal else ", no thermal constants"),
            f"grid: {g.nx} x {g.ny} nodes at {g.hx * 1e3:.2f} x {g.hy * 1e3:.2f} mm, {g.count} in the plate",
            f"thickness {hmin * 1e3:.2f} to {hmax * 1e3:.2f} mm (mean {self.mean_thickness * 1e3:.2f}); "
            f"grid dispersion -3% at {self.dispersion_frequency(0.03):.0f} Hz, -10% at {self.dispersion_frequency(0.10):.0f} Hz",
        ]
        if self.is_flat:
            lines.append("flat (no profile)")
        else:
            p2p, slope = self.profile_stats
            lines.append(f"profile: {p2p * 1e3:.1f} mm peak to peak, largest slope {slope:.3f}" + (" (above the shallow-shell range of about 0.1)" if slope > 0.1 else ""))
        geom = plate_geometry(self)
        t60 = [LN1000 / max(estimated_rate(self, geom, f).total, 1e-12) for f in (100, 1000, 5000)]
        lines.append(
            f"coincidence frequency {coincidence_frequency(m, self.mean_thickness):.0f} Hz; estimated T60 {t60[0]:.1f} / {t60[1]:.1f} / {t60[2]:.1f} s at 100 Hz / 1 kHz / 5 kHz"
        )
        return "\n".join(lines)

    def __repr__(self) -> str:
        return f"Plate({self.name!r}, {self.grid.nx}x{self.grid.ny} grid, {self.grid.count} nodes, mass {self.mass * 1e3:.0f} g)"


def remove_plane(f: np.ndarray, g: Grid) -> np.ndarray:
    """Subtracts the weighted least-squares plane a + b x + c y from ``f`` over the plate nodes."""
    w = g.node_weight
    sel = w > 0
    if not sel.any():
        return f.copy()
    A = np.stack([np.ones(sel.sum()), g.X[sel], g.Y[sel]], axis=1)
    sw = np.sqrt(w[sel])
    coef, *_ = np.linalg.lstsq(A * sw[:, None], f[sel] * sw, rcond=None)
    return f - (coef[0] + coef[1] * g.X + coef[2] * g.Y)