"""Shared pieces of the appendix diagnostics: the inputs of one system, exact spectral transforms of coarse states,
and the joint extraction of the pool3 (cFID) and spatial (csFID) Inception features."""
from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np

from bifurcation.config import RECORDS, SYSTEMS, data_paths, dataset_card, evaluation_settings, records

METHODS = ("direct_si", "edm", "cascaded_si")
LABELS = {"direct_si": "Direct SI", "edm": "EDM", "cascaded_si": "Cascaded SI"}


def add_system_args(ap: argparse.ArgumentParser, samples: bool = True) -> None:
    ap.add_argument("--system", required=True, choices=SYSTEMS)
    ap.add_argument("--data", default=None, help="dataset package directory (default: $<SYSTEM>_DATA)")
    ap.add_argument("--truth-bank", default=None, help="truth-bank directory (default: $<SYSTEM>_TRUTH_BANK)")
    if samples:
        ap.add_argument("--samples", action="append", default=[], metavar="METHOD=DIR",
                        help="test samples of a Table 2 model (direct_si, edm, cascaded_si); repeatable")
    ap.add_argument("--out", type=Path, required=True)
    ap.add_argument("--device", default="cuda")


class SystemInputs:
    """The package, truth bank, image adapter, normalization and recorded models of one system."""

    def __init__(self, a):
        from bifurcation.fid.adapters import FIDAdapter
        self.system = a.system
        self.data, self.bank = data_paths(a.system, a.data, a.truth_bank)
        self.card = dataset_card(a.system)
        self.rec = records(a.system)
        self.fid_cfg = evaluation_settings()["fid"]
        self.adapter = FIDAdapter.load(RECORDS / a.system / "fid_adapter.json", **(self.card["fid_adapter"].get("params") or {}))
        z = np.load(self.data / "normalization.npz")
        self.nmean, self.nstd = np.asarray(z["mean"], np.float64), np.asarray(z["std"], np.float64)
        self.t_idx = int(self.rec["fine_horizon_idx"])
        self.gids = sorted(int(g) for g in self.rec["test_gids"])
        self.samples = {}
        for s in getattr(a, "samples", []) or []:
            m, _, d = s.partition("=")
            if m not in METHODS or not d:
                raise SystemExit(f"--samples expects METHOD=DIR with METHOD in {METHODS}")
            self.samples[m] = Path(d).resolve()

    def unit(self, method: str) -> dict:
        key = self.rec["table2"][method]
        return next(u for u in self.rec["units"] if u["key"] == key)

    def generated(self, method: str, gid: int, n: int) -> np.ndarray:
        """RAW generated states (n, D) of test condition ``gid`` from the samples of ``method``."""
        u = self.unit(method)
        d = self.samples.get(method)
        if d is None and method == "cascaded_si" and self.rec["table2"]["cascaded_si"] == self.rec["table2"]["direct_si"]:
            d = self.samples.get("direct_si")           # S* = 1: the same model and samples as Direct SI
        if d is None:
            raise SystemExit(f"--samples {method}=DIR is needed")
        p = d / u["sampler"]["variant"] / f"group_{gid:04d}.npz"
        if not p.is_file():
            raise FileNotFoundError(f"{method}: no samples of condition {gid} at {p}")
        with np.load(p) as z:
            return np.asarray(z["gen_xT_norm"], np.float64)[:n] * self.nstd + self.nmean

    def models(self) -> list[dict]:
        """The distinct Table 2 models with samples, in the order direct_si, edm, cascaded_si (a shared model once)."""
        out, seen = [], {}
        for m in METHODS:
            key = self.rec["table2"][m]
            if key in seen:
                seen[key]["methods"].append(m); continue
            if m not in self.samples and not (m == "cascaded_si" and key == self.rec["table2"]["direct_si"] and "direct_si" in self.samples):
                continue
            seen[key] = {"key": key, "methods": [m], "n_steps": int(self.unit(m)["sampler"]["n_steps"])}
            out.append(seen[key])
        for m in out:
            m["label"] = "/".join(LABELS[x] for x in m["methods"])
        return out


# ── exact transforms of coarse spectral states (no interpolation) ──────────────────────────────────────────────

def shifted(coeff_raw, kx, ky, dx, dy, n_disp):
    """Periodic translation by (dx, dy) pixels of the n_disp grid: a phase factor on every mode."""
    M = len(kx)
    c = coeff_raw[:, :M] + 1j * coeff_raw[:, M:]
    ph = np.exp(-2j * np.pi * (kx * dx + ky * dy) / n_disp)
    c = c * ph
    return np.concatenate([c.real, c.imag], axis=1)


def rot90(coeff_raw, kx, ky, ix):
    """Rotation by ix * 90 degrees: a permutation of the Hermitian-half modes (conjugated where R k leaves the half)."""
    M = len(kx)
    c = coeff_raw[:, :M] + 1j * coeff_raw[:, M:]
    key = {(int(a), int(b)): i for i, (a, b) in enumerate(zip(kx, ky))}
    out = np.zeros_like(c)
    for i, (a, b) in enumerate(zip(kx, ky)):
        for _ in range(ix):
            a, b = -b, a
        if (int(a), int(b)) in key:
            out[:, key[(int(a), int(b))]] = c[:, i]
        elif (-int(a), -int(b)) in key:
            out[:, key[(-int(a), -int(b))]] = np.conj(c[:, i])
        else:
            raise KeyError((a, b))
    return np.concatenate([out.real, out.imag], axis=1)


def flip(coeff_raw, kx, ky, axis):
    """Reflection k_x -> -k_x (axis 0) or k_y -> -k_y (axis 1), with the Hermitian partner conjugated."""
    M = len(kx)
    c = coeff_raw[:, :M] + 1j * coeff_raw[:, M:]
    key = {(int(a), int(b)): i for i, (a, b) in enumerate(zip(kx, ky))}
    out = np.zeros_like(c)
    for i, (a, b) in enumerate(zip(kx, ky)):
        a2, b2 = (-int(a), int(b)) if axis == 0 else (int(a), -int(b))
        if (a2, b2) in key:
            out[:, key[(a2, b2)]] = c[:, i]
        else:
            out[:, key[(-a2, -b2)]] = np.conj(c[:, i])
    return np.concatenate([out.real, out.imag], axis=1)


# ── cFID and csFID features from one forward pass ──────────────────────────────────────────────────────────────

SFID_CHANNELS = 7
SFID_SPATIAL = 17
SFID_DIM = SFID_SPATIAL * SFID_SPATIAL * SFID_CHANNELS      # 2023
CFID_DIM = 2048
CANON_MIXED5_SHAPE = (768, SFID_SPATIAL, SFID_SPATIAL)


def _stages(L):
    pre = [L.conv, L.conv_1, L.conv_2, L.pool0, L.conv_3, L.conv_4, L.pool1,
           L.mixed, L.mixed_1, L.mixed_2, L.mixed_3, L.mixed_4, L.mixed_5]
    post = [L.mixed_6, L.mixed_7, L.mixed_8, L.mixed_9, L.mixed_10, L.pool2]
    return pre, post


def verify_layer_identity(extractor, device, cpu_extractor=None) -> dict:
    """Check that the stepwise traversal reproduces the pool3 feature (bit-exactly on CPU, so the spatial tap sits on
    the cFID path) and that the tapped tensor is the 17 x 17 mixed_6/conv activation. Returns the record."""
    import torch
    L = extractor.layers
    if not hasattr(L, "mixed_5") or not hasattr(L, "mixed_6") or not hasattr(L.mixed_6, "conv"):
        raise RuntimeError("this Inception build does not expose layers.mixed_6.conv")
    pre, post = _stages(L)
    x = (torch.rand(3, 3, 299, 299, device=device) * 255.0 - 128) / 128
    with torch.no_grad():
        h = x
        for st in pre:
            h = st.forward(h)
        mixed5 = h
        spatial = L.mixed_6.conv.forward(mixed5)
        for st in post:
            h = st.forward(h)
        gpu_diff = float((h.view(3, CFID_DIM) - L.forward(x).view(3, CFID_DIM)).abs().max())
    cpu_diff = None
    if cpu_extractor is not None:
        Lc = cpu_extractor.layers
        pre_c, post_c = _stages(Lc)
        xc = (torch.rand(3, 3, 299, 299) * 255.0 - 128) / 128
        with torch.no_grad():
            hc = xc
            for st in pre_c:
                hc = st.forward(hc)
            for st in post_c:
                hc = st.forward(hc)
            cpu_diff = float((hc.view(3, CFID_DIM) - Lc.forward(xc).view(3, CFID_DIM)).abs().max())
    flat = spatial[:, :SFID_CHANNELS].permute(0, 2, 3, 1).reshape(3, -1)
    rec = {"submodule": "layers.mixed_6.conv", "mixed_5_output_shape_CHW": list(mixed5.shape[1:]),
           "mixed_6_conv_output_shape_CHW": list(spatial.shape[1:]), "channels_total": int(spatial.shape[1]),
           "channels_kept": SFID_CHANNELS, "spatial_kept": [int(spatial.shape[2]), int(spatial.shape[3])],
           "feature_dim": int(flat.shape[1]), "flatten_order": "NHWC",
           "split_traversal_matches_layers_forward_bitexact_cpu": bool(cpu_diff == 0.0) if cpu_diff is not None else None,
           "split_traversal_max_abs_diff_cpu": cpu_diff, "split_traversal_max_abs_diff_cuda": gpu_diff}
    problems = []
    if tuple(mixed5.shape[1:]) != CANON_MIXED5_SHAPE:
        problems.append(f"mixed_5 output {tuple(mixed5.shape[1:])} != {CANON_MIXED5_SHAPE}")
    if tuple(spatial.shape[2:]) != (SFID_SPATIAL, SFID_SPATIAL):
        problems.append(f"mixed_6/conv spatial {tuple(spatial.shape[2:])} != {SFID_SPATIAL}x{SFID_SPATIAL}")
    if int(flat.shape[1]) != SFID_DIM:
        problems.append(f"feature dim {int(flat.shape[1])} != {SFID_DIM}")
    if cpu_diff is not None and cpu_diff != 0.0:
        problems.append(f"CPU split traversal differs from layers.forward by {cpu_diff}")
    if gpu_diff > 1e-2:
        problems.append(f"device split traversal differs from layers.forward by {gpu_diff}")
    rec["problems"] = problems
    if problems:
        raise RuntimeError("spatial feature verification failed: " + "; ".join(problems))
    print(f"Inception features verified: mixed_6/conv {list(spatial.shape[1:])}, {SFID_CHANNELS} channels kept -> "
          f"{int(flat.shape[1])}-D; stepwise traversal vs layers.forward: cpu {cpu_diff}, device {gpu_diff}\n")
    return rec


def both_features(imgs_u8, extractor, device, batch: int = 200):
    """uint8 (N, H, W, 3) -> (pool3 features (N, 2048), spatial features (N, 2023)) from one forward pass; the same
    resize and normalization as bifurcation.fid.backend."""
    import torch
    from cleanfid.resize import build_resizer
    imgs_u8 = np.asarray(imgs_u8)
    if imgs_u8.dtype != np.uint8 or imgs_u8.ndim != 4 or imgs_u8.shape[-1] != 3:
        raise ValueError(f"adapter output must be uint8 (N, H, W, 3); got {imgs_u8.dtype} {imgs_u8.shape}")
    L = extractor.layers
    pre, _ = _stages(L)
    resizer = build_resizer("clean")
    fc, fs = [], []
    with torch.no_grad():
        for i in range(0, len(imgs_u8), int(batch)):
            xs = [resizer(im) for im in imgs_u8[i:i + int(batch)]]
            x = torch.from_numpy(np.stack(xs).transpose(0, 3, 1, 2).copy()).to(device)
            if tuple(x.shape[-2:]) != (299, 299):
                raise AssertionError("resizer did not produce 299x299")
            fc.append(extractor(x).cpu().numpy())
            h = (x - 128) / 128
            for st in pre:
                h = st.forward(h)
            sp = L.mixed_6.conv.forward(h)[:, :SFID_CHANNELS]
            fs.append(sp.permute(0, 2, 3, 1).reshape(sp.shape[0], -1).cpu().numpy())
    return np.concatenate(fc).astype(np.float32), np.concatenate(fs).astype(np.float32)


def extractors(fid_cfg: dict, device: str):
    """(extractor on the device, the layer-identity record); the pins of the fid settings are checked first."""
    import torch
    from bifurcation.fid import backend as BE
    BE.check_pins(fid_cfg)
    dev = torch.device(device if torch.cuda.is_available() else "cpu")
    ext = BE.build_extractor(dev)
    cpu_ext = BE.build_extractor(torch.device("cpu")) if dev.type != "cpu" else ext
    rec = verify_layer_identity(ext, dev, cpu_extractor=cpu_ext)
    return ext, dev, rec


def write_json(path: Path, obj) -> None:
    path.write_text(json.dumps(obj, indent=1, default=float) + "\n")