"""The modes of a plate up to a cutoff frequency, and their loss rates.

The generalized problem K phi = omega^2 M phi with the diagonal lumped mass M
is turned into the standard symmetric problem A = M^{-1/2} K M^{-1/2}, whose
eigenvalues in the band of interest are computed by LAPACK, on the CPU through
SciPy (a dense solve restricted to the band; about 4 s for a 16 inch cymbal at
8 mm, and 40 s for the same cymbal with a profile, whose system is three times
as large) or on a GPU through PyTorch (a full dense eigendecomposition in
double precision, the faster path for the larger systems when a CUDA device
is available). The mode shapes
returned are M^{-1/2} times the eigenvectors of A, that is, mass-normalized.
The zero-frequency rigid modes of the free plate are dropped.
"""

from __future__ import annotations

import math
import time
from dataclasses import dataclass, replace
from typing import Literal

import numpy as np
import scipy.linalg

from ..constants import LN1000
from ..plate import Plate
from .assemble import LinearSystem, assemble

Solver = Literal["auto", "banded", "dense", "torch"]


@dataclass
class Modes:
    """The modes of a plate below a cutoff frequency, with their per-mode decay rates.

    Shapes are mass-normalized (phi^T M phi = 1) and given on the grid, zero
    outside the plate. For a flat plate only the transverse component ``w``
    exists; a plate with a profile also has the in-plane components ``u`` and ``v``.
    """

    plate: Plate
    freqs: np.ndarray
    """Frequencies [Hz], ascending, ``(M,)``."""
    w: np.ndarray
    """Transverse component of the shapes, ``(M, ny, nx)``."""
    u: np.ndarray | None
    v: np.ndarray | None
    sigma_material: np.ndarray
    """Decay rate of each mode from internal friction [1/s]."""
    sigma_thermal: np.ndarray
    """Decay rate from thermoelastic damping [1/s]."""
    sigma_radiation: np.ndarray
    """Decay rate from radiation into the air [1/s]."""
    fmax: float
    """Cutoff frequency of the computation [Hz]."""
    seconds: float = 0.0
    """Wall time of the computation [s]."""

    @property
    def count(self) -> int:
        return len(self.freqs)

    @property
    def omega(self) -> np.ndarray:
        return 2 * math.pi * self.freqs

    @property
    def sigma(self) -> np.ndarray:
        """Total decay rate of each mode [1/s], the sum of the three mechanisms, kept below the mode's frequency (underdamped)."""
        return np.minimum(self.sigma_material + self.sigma_thermal + self.sigma_radiation, 0.99 * self.omega)

    @property
    def t60(self) -> np.ndarray:
        """Decay time of each mode [s]: the time in which it falls by 60 dB."""
        with np.errstate(divide="ignore"):
            return LN1000 / self.sigma

    def t60_by_mechanism(self) -> dict[str, np.ndarray]:
        """The decay time each mechanism alone would give: ``material``, ``thermal``, ``radiation``."""
        out = {}
        for name, s in (("material", self.sigma_material), ("thermal", self.sigma_thermal), ("radiation", self.sigma_radiation)):
            with np.errstate(divide="ignore"):
                out[name] = LN1000 / s
        return out

    def select(self, index) -> "Modes":
        """The modes at the given indices (an integer array, slice or boolean mask)."""
        idx = np.arange(self.count)[index]
        return replace(
            self,
            freqs=self.freqs[idx],
            w=self.w[idx],
            u=None if self.u is None else self.u[idx],
            v=None if self.v is None else self.v[idx],
            sigma_material=self.sigma_material[idx],
            sigma_thermal=self.sigma_thermal[idx],
            sigma_radiation=self.sigma_radiation[idx],
        )

    def below(self, fmax: float) -> "Modes":
        """The modes below a frequency [Hz]."""
        return self.select(self.freqs < fmax)

    def table(self, n: int | None = None) -> str:
        """A text table of the first ``n`` modes: frequency and decay times (total, and by mechanism)."""
        by = self.t60_by_mechanism()
        rows = ["  #    f [Hz]   T60 [s]   air   metal  thermal"]
        for m in range(self.count if n is None else min(n, self.count)):
            rows.append(f"{m:3d} {self.freqs[m]:9.1f} {self.t60[m]:9.2f} {_fmt(by['radiation'][m])} {_fmt(by['material'][m])} {_fmt(by['thermal'][m])}")
        return "\n".join(rows)

    def save(self, path) -> None:
        """Saves the modes and the plate to an ``.npz`` file (see ``Modes.load``)."""
        from ..material import Thermal

        pl, m = self.plate, self.plate.material
        th = m.thermal or Thermal(0, 1, 1)
        np.savez_compressed(
            path,
            freqs=self.freqs,
            w=self.w,
            u=np.zeros(0) if self.u is None else self.u,
            v=np.zeros(0) if self.v is None else self.v,
            sigma_material=self.sigma_material,
            sigma_thermal=self.sigma_thermal,
            sigma_radiation=self.sigma_radiation,
            fmax=self.fmax,
            seconds=self.seconds,
            mask=pl.grid.mask,
            hx=pl.grid.hx,
            hy=pl.grid.hy,
            h=pl.h,
            w0=pl.w0,
            Lx=pl.Lx,
            Ly=pl.Ly,
            temperature=pl.temperature,
            name=pl.name,
            material=np.array([m.E, m.nu, m.rho, m.loss_factor, 1.0 if m.thermal else 0.0, th.alpha, th.cp, th.k]),
        )

    @classmethod
    def load(cls, path) -> "Modes":
        from ..grid import Grid
        from ..material import Material, Thermal

        d = np.load(path)
        mat = d["material"]
        material = Material(float(mat[0]), float(mat[1]), float(mat[2]), float(mat[3]), Thermal(float(mat[5]), float(mat[6]), float(mat[7])) if mat[4] > 0 else None)
        mask = d["mask"]
        grid = Grid(mask.shape[1], mask.shape[0], float(d["hx"]), float(d["hy"]), mask)
        plate = Plate(grid, material, d["h"], d["w0"], float(d["Lx"]), float(d["Ly"]), float(d["temperature"]), str(d["name"]))
        u, v = d["u"], d["v"]
        return cls(plate, d["freqs"], d["w"], None if u.size == 0 else u, None if v.size == 0 else v, d["sigma_material"], d["sigma_thermal"], d["sigma_radiation"], float(d["fmax"]), float(d["seconds"]))

    def __repr__(self) -> str:
        return f"Modes({self.count} modes of {self.plate.name!r} up to {self.fmax:.0f} Hz)"


def _fmt(t: float) -> str:
    if not np.isfinite(t) or t <= 0:
        return "    inf"
    if t < 10:
        return f"{t:7.2f}"
    if t < 1000:
        return f"{t:7.0f}"
    return f"{t:7.0e}"


def compute_modes(
    plate: Plate,
    fmax: float = 16000.0,
    *,
    fmin: float = 1.0,
    solver: Solver = "auto",
    device=None,
    damping: bool = True,
    system: LinearSystem | None = None,
) -> Modes:
    """All modes of the plate with frequencies in ``(fmin, fmax]``.

    ``solver`` picks the eigensolver: ``"dense"`` (SciPy on the full matrix,
    eigenvalues in the band only; the default on the CPU), ``"banded"``
    (LAPACK's banded driver, which uses less memory but is slower here),
    ``"torch"`` (a dense eigendecomposition on ``device``, useful on a CUDA
    GPU), or ``"auto"``, which takes the torch path when a CUDA device is
    available and the dense path otherwise. With ``damping`` the per-mode
    thermoelastic and radiation rates are computed from the shapes (a second or
    two for a few thousand modes); without it they are left at zero.
    """
    t0 = time.perf_counter()
    sys = system or assemble(plate)
    lo = (2 * math.pi * fmin) ** 2
    hi = (2 * math.pi * fmax) ** 2
    if solver == "auto":
        solver = "torch" if _cuda_available() else "dense"
    if solver == "banded":
        lam, y = _eig_banded(sys, lo, hi)
    elif solver == "dense":
        lam, y = _eig_dense(sys, lo, hi)
    elif solver == "torch":
        lam, y = _eig_torch(sys, lo, hi, device)
    else:
        raise ValueError(f"unknown solver {solver!r}")
    order = np.argsort(lam)
    lam, y = lam[order], y[:, order]
    phi = (y / np.sqrt(sys.mass)[:, None]).T  # (M, n), mass-normalized
    u, v, w = sys.to_grid(phi)
    freqs = np.sqrt(np.maximum(lam, 0)) / (2 * math.pi)
    M = len(freqs)
    modes = Modes(
        plate,
        freqs,
        w,
        u if sys.fields == 3 else None,
        v if sys.fields == 3 else None,
        plate.material.loss_factor * (2 * math.pi * freqs) / 2,
        np.zeros(M),
        np.zeros(M),
        fmax,
    )
    if damping and M > 0:
        from .damping import radiation_rates, thermoelastic_rates

        modes.sigma_radiation = radiation_rates(modes)
        modes.sigma_thermal = thermoelastic_rates(modes)
    modes.seconds = time.perf_counter() - t0
    return modes


def _cuda_available() -> bool:
    try:
        import torch

        return torch.cuda.is_available()
    except ImportError:
        return False


def _scaled(sys: LinearSystem):
    s = 1 / np.sqrt(sys.mass)
    A = sys.K.multiply(s[:, None]).multiply(s[None, :]).tocoo()
    return A, s


def _eig_banded(sys: LinearSystem, lo: float, hi: float):
    A, _ = _scaled(sys)
    sel = A.row >= A.col
    r, c, v = A.row[sel], A.col[sel], A.data[sel]
    bw = int((r - c).max()) if len(r) else 0
    ab = np.zeros((bw + 1, sys.n))
    ab[r - c, c] = v
    lam, y = scipy.linalg.eig_banded(ab, lower=True, select="v", select_range=(lo, hi), check_finite=False)
    return lam, y


def _eig_dense(sys: LinearSystem, lo: float, hi: float):
    A, _ = _scaled(sys)
    lam, y = scipy.linalg.eigh(A.toarray(), subset_by_value=(lo, hi), driver="evr", check_finite=False)
    return lam, y


def _eig_torch(sys: LinearSystem, lo: float, hi: float, device):
    import torch

    A, _ = _scaled(sys)
    dev = torch.device(device) if device is not None else torch.device("cuda" if torch.cuda.is_available() else "cpu")
    At = torch.from_numpy(A.toarray()).to(dev, dtype=torch.float64)
    lam, y = torch.linalg.eigh(At)
    lam = lam.cpu().numpy()
    sel = (lam > lo) & (lam <= hi)
    return lam[sel], y[:, torch.from_numpy(sel).to(dev)].cpu().numpy()