"""NLKG image adapter: score displacement while tracking both displacement and velocity diagnostics.

The 392-dimensional state contains 196 displacement coordinates followed by 196 velocity
coordinates scaled by 1/4. Images reconstruct the displacement on a 128 x 128 grid.
"""
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

_D_C = 196


@register_adapter("nlkg_coarse_u")
class NLKGCoarseUAdapter(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
        if 2 * self.n_modes != _D_C:
            raise ValueError("nlkg_coarse_u assumes the 196-D shared coarse operator per component")

    @staticmethod
    def _split(x: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        x = np.asarray(x, dtype=np.float64)
        if x.ndim == 1:
            x = x[None]
        if x.shape[1] == 2 * _D_C:
            return x[:, :_D_C], x[:, _D_C:]
        if x.shape[1] == _D_C:
            return x, None
        raise ValueError(f"expected {2 * _D_C}-D two-component (u, v/4) states or {_D_C}-D displacement-only states, got D={x.shape[1]}")

    def to_fields(self, states_raw: np.ndarray) -> np.ndarray:
        re, _ = self._split(states_raw)
        return super().to_fields(re)

    def card(self) -> dict:
        c = super().card()
        c.update({"state": "392-D coarse state (196-D u + 196-D v/4, shared coarse operator |k| <= 8 on each); "
                           "the image reads only the first component u -- identical reconstruction to euler_coarse_vorticity",
                  "clip": "symmetric [-c, c], c = Q_q(|u_coarse|) over the training pool's final-time fields",
                  "physical_diagnostic": "coarse_u_variance_ratio, gradient_energy_ratio (on u), coarse_v_variance_ratio (on v/4)"})
        return c


    def _power(self, v: np.ndarray) -> tuple[float, float]:
        n = self.n_modes
        power = v[:, :n] ** 2 + v[:, n:] ** 2
        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:
        gr, gi = self._split(gen_raw); tr, ti = self._split(true_raw)
        vg, eg = self._power(gr); vt, et = self._power(tr)
        out = {"coarse_u_variance_ratio": vg / vt if vt > 0 else float("nan"),
               "gradient_energy_ratio": eg / et if et > 0 else float("nan")}
        if gi is not None and ti is not None:
            vgi, _ = self._power(gi); vti, _ = self._power(ti)
            out["coarse_v_variance_ratio"] = vgi / vti if vti > 0 else float("nan")
        return out