"""Heterogeneous-wave FID adapter: the SAME 196-D coarse spectral observable as Euler and the passive scalar (98 complex modes |k| <= 8, Hermitian half, coefficients / N^2, shared mode order from
``bifurcation.models.spectral``) -> the exact low-pass DISPLACEMENT field u on a 128x128 grid -> symmetric clip ->
uint8 -> 3 identical channels.

The reconstruction is byte-identical to the Euler adapter's (one inverse FFT, one mode-index map); only the physical
reading differs. The observed field is the wave displacement u of

    u_t = v,    v_t = div(c^2 grad u) - gamma v,    c^2(x, y) = 1 + 0.5 sin(2 pi x) sin(2 pi y),

which is signed and zero-mean (the DC mode is not retained by the observable), so the symmetric clip about zero is
the right quantization (a nonnegative [0, c] range would throw away half of every wave crest-trough pair).

The velocity v and the unresolved band 8 < |k| <= 16 are hidden, which is what makes p(x_T | x_0) non-degenerate;
neither enters the image, which is a function of the observed state alone.

Physical diagnostics, both computed from the coefficient vectors:
  * ``displacement_variance_ratio`` — mean sum |u_hat_k|^2, the direct analogue of Euler's enstrophy ratio and the
    passive scalar's scalar-variance ratio;
  * ``gradient_energy_ratio`` — mean sum |k|^2 |u_hat_k|^2, the resolved part of the potential energy density
    c^2 |grad u|^2 up to the medium factor; a wave field can match the first moment while carrying its energy at
    the wrong wavenumber, and this separates the two.
"""
from __future__ import annotations

import numpy as np

from ..models.spectral import K_C, mode_int_wavenumbers
from .adapters import register_adapter
from .euler import EulerVorticityAdapter


@register_adapter("heterogeneous_waves_coarse")
class HeterogeneousWavesAdapter(EulerVorticityAdapter):
    field_kind = "signed"

    def __init__(self, quantile: float = 0.9999, n_disp: int = 128, k_c: int = K_C, **kwargs):
        super().__init__(quantile, n_disp=n_disp, k_c=k_c, **kwargs)
        kx, ky = mode_int_wavenumbers(self.k_c)
        self._k2 = kx.astype(np.float64) ** 2 + ky.astype(np.float64) ** 2

    def card(self) -> dict:
        c = super().card()
        c.update({"state": "196-D coarse spectral wave displacement (98 complex modes |k| <= 8, Hermitian half, u_hat / N^2)",
                  "clip": "symmetric [-c, c], c = Q_q(|u|) over the training pool",
                  "physical_diagnostic": "displacement_variance_ratio = mean_gen sum|u_hat_k|^2 / mean_truth sum|u_hat_k|^2; "
                                         "gradient_energy_ratio = mean_gen sum|k|^2|u_hat_k|^2 / mean_truth sum|k|^2|u_hat_k|^2"})
        return c

    def _moments(self, v: np.ndarray) -> tuple[float, float]:
        v = np.asarray(v, np.float64)
        n = self.n_modes
        power = v[:, :n] ** 2 + v[:, n:] ** 2          # |u_hat_k|^2 per retained mode
        return float(np.mean(np.sum(power, axis=1))), float(np.mean(np.sum(self._k2 * power, axis=1)))

    def physical_diagnostics(self, gen_raw: np.ndarray, true_raw: np.ndarray) -> dict:
        dg, gg = self._moments(gen_raw)
        dt, gt = self._moments(true_raw)
        return {"displacement_variance_ratio": dg / dt if dt > 0 else float("nan"),
                "gradient_energy_ratio": gg / gt if gt > 0 else float("nan")}