"""Coarse spectral <-> field maps for the 196-dim state (numpy, no plotting).

The (kx, ky) <-> index map is the single source of truth for the 196-dim coarse spectral vector shared by the
five datasets.

Reconstruction (exact inverse of extract_group.get_coarse_vec):
  vec = [Re(oh_norm[mask]), Im(oh_norm[mask])], oh_norm = omega_hat / N_sim^2,
  mask = (|k| <= k_c) & Hermitian-half. Un-normalize (clouds are per-feature
  standardized), place each complex coefficient and its conjugate partner on an
  N_disp grid, field = N_disp^2 * ifft2(OH).real.

There is deliberately no default dataset path in this module: the normalization
statistics that the wrapper un/re-normalizes with must come from the dataset
the model is trained on (the registry injects ``data.packaged_dataset_path``).
"""

from __future__ import annotations

from pathlib import Path

import numpy as np

K_C = 8
N_SIM = 1024  # the resolution get_coarse_vec normalized by (omega_hat / N_sim^2)


def mode_int_wavenumbers(k_c: int = K_C, M: int = N_SIM):
    """Integer (kx, ky) for each of the 98 complex modes, in get_coarse_vec
    order. axis0<->kx, axis1<->ky (matches extract_group.make_kgrids)."""
    k1d = np.fft.fftfreq(M, d=1.0 / M).astype(int)      # 0,1,..,-1
    ky_grid, kx_grid = np.meshgrid(k1d, k1d)
    k_mag = np.sqrt(kx_grid**2 + ky_grid**2)
    hermitian = (kx_grid > 0) | ((kx_grid == 0) & (ky_grid > 0))
    mask = (k_mag <= k_c) & hermitian
    return kx_grid[mask], ky_grid[mask]                 # each (n_modes,)


class FieldReconstructor:
    def __init__(self, packaged: str | Path,
                 n_disp: int = 128, k_c: int = K_C):
        norm = np.load(Path(packaged) / "normalization.npz")
        self.mean, self.std = norm["mean"], norm["std"]
        self.n_disp = n_disp
        self.kx, self.ky = mode_int_wavenumbers(k_c)
        self.n_modes = len(self.kx)
        # precompute grid placement indices (mod N for negative wavenumbers)
        self.ix = self.kx % n_disp
        self.iy = self.ky % n_disp
        self.ix_c = (-self.kx) % n_disp
        self.iy_c = (-self.ky) % n_disp

    def _spectral_grid(self, vec_norm: np.ndarray) -> np.ndarray:
        """(196,) standardized coarse vector -> full Hermitian ω̂ grid (n_disp²)."""
        vec = np.asarray(vec_norm, float) * self.std + self.mean   # un-normalize
        coeff = vec[:self.n_modes] + 1j * vec[self.n_modes:]       # oh_norm[mask]
        OH = np.zeros((self.n_disp, self.n_disp), dtype=complex)
        OH[self.ix,   self.iy]   += coeff
        OH[self.ix_c, self.iy_c] += np.conj(coeff)
        return OH

    def reconstruct_field(self, vec_norm: np.ndarray,
                          field: str = "vorticity") -> np.ndarray:
        """(196,) standardized coarse vector -> (N_disp, N_disp) real field.

        Only the 'vorticity' branch (the observed scalar itself) exists.
        """
        OH = self._spectral_grid(vec_norm)
        n2 = self.n_disp ** 2
        if field == "vorticity":
            return n2 * np.fft.ifft2(OH).real
        raise ValueError(f"unknown field: {field!r} (only 'vorticity' is supported here)")

    def reconstruct_many(self, vecs_norm: np.ndarray,
                         field: str = "vorticity") -> np.ndarray:
        """(N, 196) -> (N, N_disp, N_disp)."""
        return np.stack([self.reconstruct_field(v, field) for v in vecs_norm])

    def extract_vec(self, field: np.ndarray) -> np.ndarray:
        """(N_disp, N_disp) real field -> (2*n_modes,) standardized coarse vector.

        Exact left-inverse of reconstruct_field when the field is band-limited
        to |k| <= k_c on this grid (no aliasing): OH = fft2(field)/N_disp^2 undoes
        the N_disp^2 scaling in reconstruct_field, the placed coefficients are
        read back at the mode indices, and the per-feature standardization is
        re-applied. Used to verify the FFT wrapper round-trips and as the forward
        leg of the unet2 spectral<->field map."""
        OH = np.fft.fft2(np.asarray(field, float)) / (self.n_disp ** 2)
        coeff = OH[self.ix, self.iy]
        vec = np.concatenate([coeff.real, coeff.imag])
        return (vec - self.mean) / self.std

    def roundtrip_maxerr(self, vecs_norm: np.ndarray) -> float:
        """max |extract_vec(reconstruct_field(v)) - v| over the given vectors."""
        vecs_norm = np.atleast_2d(np.asarray(vecs_norm, float))
        errs = [np.abs(self.extract_vec(self.reconstruct_field(v)) - v).max()
                for v in vecs_norm]
        return float(max(errs))