"""Dataset adapters psi_D: physical state -> three-channel uint8 image.

Contract (enforced by the base class):
  * ``fit(train_states_raw)`` computes the adapter's numerical constants from training data only and fixes
    them (``params``; records/<system>/fid_adapter.json);
  * ``to_images(states_raw)`` is deterministic given ``params``; no per-image normalization;
  * signed fields with a meaningful zero: symmetric clip [-c, c], c = Q_q(|x_train|); nonnegative fields:
    [0, c], c = Q_q(x_train); the map to uint8 is the fixed linear map of the range onto [0, 255];
  * scalar fields -> three identical channels unless the dataset card defines a physical 3-channel construction.

States enter raw (physical units, un-standardized). configs/<system>/dataset.yaml names the adapter and its parameters.
"""
from __future__ import annotations

import json
from pathlib import Path

import numpy as np

ADAPTERS: dict[str, type] = {}


def register_adapter(name: str):
    def deco(cls):
        cls.name = name
        ADAPTERS[name] = cls
        return cls
    return deco


def get_adapter(name: str, **kwargs) -> "FIDAdapter":
    if name not in ADAPTERS:
        raise KeyError(f"unknown FID adapter {name!r} (registered: {sorted(ADAPTERS)})")
    return ADAPTERS[name](**kwargs)


def quantize_symmetric(fields: np.ndarray, c: float) -> tuple[np.ndarray, float]:
    """Signed field -> uint8 grayscale: rint((clip(f, -c, c) + c) / (2c) * 255); zero -> 128. Returns the
    clipped-pixel fraction too."""
    f = np.asarray(fields, dtype=np.float64)
    clipped = float(np.mean((f > c) | (f < -c))) if f.size else 0.0
    g = np.rint((np.clip(f, -c, c) + c) / (2.0 * c) * 255.0)
    return np.clip(g, 0, 255).astype(np.uint8), clipped


def quantize_nonnegative(fields: np.ndarray, c: float) -> tuple[np.ndarray, float]:
    f = np.asarray(fields, dtype=np.float64)
    clipped = float(np.mean(f > c)) if f.size else 0.0
    g = np.rint(np.clip(f, 0.0, c) / c * 255.0)
    return np.clip(g, 0, 255).astype(np.uint8), clipped


def gray_to_rgb(g: np.ndarray) -> np.ndarray:
    return np.repeat(g[..., None], 3, axis=-1)


class FIDAdapter:
    name = "abstract"
    field_kind = "signed"           # 'signed' (symmetric clip) | 'nonnegative' ([0, c])

    def __init__(self, quantile: float = 0.9999, **kwargs):
        self.quantile = float(quantile)
        self.params: dict | None = None
        self.kwargs = dict(kwargs)

    # dataset-specific: raw state vectors (N, D) -> scalar fields (N, H, W)
    def to_fields(self, states_raw: np.ndarray) -> np.ndarray:
        raise NotImplementedError

    def fit(self, train_states_raw: np.ndarray, source: dict | None = None, pixel_stride: int = 1,
            chunk: int = 4000) -> dict:
        """Constants from TRAINING data only. Freezes ``self.params``. ``pixel_stride`` subsamples the field
        pixels (every k-th in each axis);
        states are converted in chunks to bound memory."""
        k = int(pixel_stride)
        vals, shape = [], None
        for i in range(0, len(train_states_raw), int(chunk)):
            f = self.to_fields(train_states_raw[i:i + int(chunk)])
            shape = list(f.shape[1:])
            f = f[:, ::k, ::k] if f.ndim == 3 else f
            vals.append((np.abs(f) if self.field_kind == "signed" else f).astype(np.float32).ravel())
        vals = np.concatenate(vals)
        c = float(np.quantile(vals, self.quantile))
        rng = [-c, c] if self.field_kind == "signed" else [0.0, c]
        self.params = {"adapter": self.name, "field_kind": self.field_kind, "quantile": self.quantile, "clip_c": c,
                       "range": rng, "n_train_states": int(len(train_states_raw)), "field_shape": shape,
                       "pixel_stride": k, "n_pixel_values": int(vals.size),
                       "map": ("uint8 = rint((clip(w,-c,c)+c)/(2c)*255)" if self.field_kind == "signed"
                               else "uint8 = rint(clip(w,0,c)/c*255)"),
                       "channels": "3 identical (grayscale scalar field), no colormap",
                       "source": source or {}, **self.card()}
        return self.params

    def card(self) -> dict:
        """Dataset-card text of the adapter (channel construction, reconstruction) — override."""
        return {}

    def load_params(self, params: dict) -> "FIDAdapter":
        if params.get("adapter") != self.name:
            raise ValueError(f"params belong to adapter {params.get('adapter')!r}, not {self.name!r}")
        self.params = dict(params)
        self.quantile = float(params.get("quantile", self.quantile))
        return self

    def to_images(self, states_raw: np.ndarray) -> tuple[np.ndarray, float]:
        """RAW states -> uint8 (N, H, W, 3) and the clipped-pixel fraction. Requires fitted params."""
        if self.params is None:
            raise RuntimeError("adapter not fitted / loaded (constants come from training data only)")
        fields = self.to_fields(states_raw)
        c = float(self.params["clip_c"])
        g, frac = (quantize_symmetric(fields, c) if self.field_kind == "signed" else quantize_nonnegative(fields, c))
        return gray_to_rgb(g), frac

    def physical_diagnostics(self, gen_raw: np.ndarray, true_raw: np.ndarray) -> dict:
        """Optional physical diagnostics of a generated ensemble against the truth. Default: none."""
        return {}

    def save(self, path: Path) -> Path:
        Path(path).write_text(json.dumps(self.params, indent=1) + "\n")
        return Path(path)

    @classmethod
    def load(cls, path: Path, **kwargs) -> "FIDAdapter":
        params = json.loads(Path(path).read_text())
        return get_adapter(params["adapter"], quantile=params.get("quantile", 0.9999), **kwargs).load_params(params)


@register_adapter("vector_grid")
class VectorGridAdapter(FIDAdapter):
    """Generic adapter for vector states with no field structure: the D coordinates are laid out row-major on
    the smallest square grid (zero padding), each cell drawn as a ``cell`` x ``cell`` block (nearest
    upsampling, no interpolation), symmetric clip from the training quantile. Used for datasets whose state is
    a plain vector and for the synthetic test dataset."""
    field_kind = "signed"

    def __init__(self, quantile: float = 0.9999, cell: int = 8, **kwargs):
        super().__init__(quantile, **kwargs)
        self.cell = int(cell)

    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]
        N, D = x.shape
        side = int(np.ceil(np.sqrt(D)))
        grid = np.zeros((N, side * side), dtype=np.float64)
        grid[:, :D] = x
        grid = grid.reshape(N, side, side)
        return np.repeat(np.repeat(grid, self.cell, axis=1), self.cell, axis=2)

    def card(self) -> dict:
        return {"reconstruction": f"row-major {int(np.ceil(np.sqrt(1)))}x? square grid of the state coordinates, zero-padded, "
                                  f"each coordinate a {self.cell}x{self.cell} block (nearest)", "cell": self.cell}