"""Euler FID adapter: 196-D coarse spectral vorticity (|k| <= 8, Hermitian half, raw coefficients
omega_hat / N_sim^2) -> the EXACT low-pass vorticity field on a 128x128 grid -> symmetric clip -> uint8 -> 3
identical channels.

The reconstruction is the same inverse FFT as ``bifurcation.models.spectral.FieldReconstructor`` (the mode
index map is imported from there — one source of truth), vectorized over N states; equality with the
reference per-vector reconstruction.

Constant: c = Q_0.9999(|omega|) over the training pool (records/euler/fid_adapter.json). Physical diagnostic:
the enstrophy ratio (mean enstrophy of the generated ensemble / mean enstrophy of the truth, from the coefficients).
"""
from __future__ import annotations

import numpy as np

from ..models.spectral import K_C, mode_int_wavenumbers
from .adapters import FIDAdapter, register_adapter


@register_adapter("euler_coarse_vorticity")
class EulerVorticityAdapter(FIDAdapter):
    field_kind = "signed"

    def __init__(self, quantile: float = 0.9999, n_disp: int = 128, k_c: int = K_C, **kwargs):
        super().__init__(quantile, **kwargs)
        self.n_disp, self.k_c = int(n_disp), int(k_c)
        kx, ky = mode_int_wavenumbers(self.k_c)
        self.n_modes = len(kx)
        self.ix, self.iy = kx % self.n_disp, ky % self.n_disp
        self.ix_c, self.iy_c = (-kx) % self.n_disp, (-ky) % self.n_disp
        if self.k_c * 2 >= self.n_disp:
            raise ValueError("n_disp must exceed 2 k_c (alias-free reconstruction)")

    def to_fields(self, states_raw: np.ndarray) -> np.ndarray:
        x = np.asarray(states_raw, dtype=np.float64)
        if x.ndim == 1:
            x = x[None]
        if x.shape[1] != 2 * self.n_modes:
            raise ValueError(f"expected {2 * self.n_modes}-D coarse spectral states, got D={x.shape[1]}")
        N = x.shape[0]
        coeff = x[:, :self.n_modes] + 1j * x[:, self.n_modes:]
        OH = np.zeros((N, self.n_disp, self.n_disp), dtype=np.complex128)
        OH[:, self.ix, self.iy] += coeff
        OH[:, self.ix_c, self.iy_c] += np.conj(coeff)
        return (self.n_disp ** 2) * np.fft.ifft2(OH, axes=(1, 2)).real

    def card(self) -> dict:
        return {"state": "196-D coarse spectral vorticity (98 complex modes |k| <= 8, Hermitian half, omega_hat / N_sim^2)",
                "reconstruction": f"exact inverse FFT of the |k| <= {self.k_c} modes on a {self.n_disp}x{self.n_disp} grid (alias-free)",
                "n_disp": self.n_disp, "k_c": self.k_c, "clip": "symmetric [-c, c], c = Q_q(|omega_T|) over the training pool's final-time fields",
                "physical_diagnostic": "enstrophy_ratio = mean_gen sum|omega_hat_k|^2 / mean_truth sum|omega_hat_k|^2"}

    def physical_diagnostics(self, gen_raw: np.ndarray, true_raw: np.ndarray) -> dict:
        def ens(v):
            v = np.asarray(v, np.float64)
            return float(np.mean(np.sum(v ** 2, axis=1)))       # sum over Re/Im of all Hermitian-half modes
        et = ens(true_raw)
        return {"enstrophy_ratio": ens(gen_raw) / et if et > 0 else float("nan")}