"""The Clean-FID pipeline, pinned and checked.

Everything after the dataset adapter is fixed here for every dataset:
  * images arrive as uint8 (N, H, W, 3) in [0, 255] (one fixed image range);
  * a single resize to 299x299 with clean-fid's "clean" resizer (PIL bicubic, float output, antialiasing is
    PIL's default for its resampling filters); the TorchScript Inception is built with resize_inside=False
    and asserts 299x299 input, so a second resize cannot happen silently;
  * the pinned Inception-v3 weights (inception-2015-12-05.pt, sha256 checked), pool3 features (2048-D);
  * FID = clean-fid's frechet_distance formula; KID = clean-fid's kernel_distance (unbiased polynomial-kernel
    MMD^2, num_subsets / max_subset_size from the evaluation settings).

``pins()`` returns the record that every FID manifest carries; ``check_pins(expected)`` raises on any drift.
"""
from __future__ import annotations

import hashlib
import importlib.metadata as _md
import os
from pathlib import Path

import numpy as np

INCEPTION_DIR = Path(os.environ.get("CLEAN_FID_CACHE", Path.home() / ".cache" / "clean-fid"))
INCEPTION_FILE = "inception-2015-12-05.pt"
IMAGE_SIZE = 299

# ── the spatial Inception feature of csFID (Nash et al., ICML 2021) ─────────────────────
# The Inception-v3 activation "mixed_6/conv" -- the FIRST branch of the mixed_6 Inception block, BEFORE its
# concatenation with tower/tower_1/tower_2 (openai/guided-diffusion evaluator.py FID_SPATIAL_NAME) -- is a
# 192 x 17 x 17 tensor; its first 7 channels are kept and flattened NHWC to 17*17*7 = 2023 components per image,
# without spatial average pooling (that pooling is what makes the 2048-D pool3 cFID feature layout-tolerant).
# In this TorchScript port the tensor is the submodule layers.mixed_6.conv; the module names come from the TF op
# names. The stepwise traversal used to read it reproduces layers.forward exactly on CPU (verify_spatial_tap).
SPATIAL_LAYER = "mixed_6/conv"
SPATIAL_CHANNELS = 7
SPATIAL_GRID = 17
SPATIAL_DIM = SPATIAL_GRID * SPATIAL_GRID * SPATIAL_CHANNELS          # 2023
POOL3_DIM = 2048
_MIXED5_SHAPE = (768, SPATIAL_GRID, SPATIAL_GRID)                     # canonical input of the mixed_6 block


def _sha256(path: Path) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 22), b""):
            h.update(chunk)
    return h.hexdigest()


def pins() -> dict:
    import PIL
    import torch
    p = INCEPTION_DIR / INCEPTION_FILE
    return {"backend": "clean-fid", "clean_fid_version": _md.version("clean-fid"), "mode": "clean",
            "resize": {"backend": "PIL", "filter": "bicubic", "size": IMAGE_SIZE, "once": True,
                       "implementation": "cleanfid.resize.build_resizer('clean') = make_resizer('PIL', False, 'bicubic', (299, 299))"},
            "inception": {"file": INCEPTION_FILE, "path": str(p), "sha256": _sha256(p) if p.is_file() else None,
                          "feature_layer": "pool3", "feature_dim": POOL3_DIM, "resize_inside": False},
            "spatial": {"feature_layer": SPATIAL_LAYER, "submodule": "layers.mixed_6.conv",
                        "tap": "first branch of the mixed_6 block, before concatenation with tower/tower_1/tower_2",
                        "channels_kept": SPATIAL_CHANNELS, "grid": [SPATIAL_GRID, SPATIAL_GRID], "feature_dim": SPATIAL_DIM,
                        "flatten": "NHWC", "spatial_pooling": "none",
                        "reference": "Nash et al., ICML 2021; openai/guided-diffusion evaluator.py FID_SPATIAL_NAME='mixed_6/conv:0'"},
            "fid": "cleanfid.fid.frechet_distance (scipy sqrtm, eps 1e-6 fallback)",
            "kid": "cleanfid.fid.kernel_distance (unbiased polynomial-kernel MMD^2)",
            "torch": torch.__version__, "pillow": PIL.__version__, "numpy": np.__version__}


def check_pins(expected: dict) -> dict:
    """``expected`` = the fid settings (clean_fid_version, inception.sha256, inception.file). Raises on drift."""
    got = pins()
    probs = []
    if str(expected.get("clean_fid_version")) != got["clean_fid_version"]:
        probs.append(f"clean-fid {got['clean_fid_version']} != pinned {expected.get('clean_fid_version')}")
    exp_inc = expected.get("inception") or {}
    if exp_inc.get("file") and exp_inc["file"] != got["inception"]["file"]:
        probs.append("inception file name differs")
    if got["inception"]["sha256"] is None:
        probs.append(f"Inception weights not found at {got['inception']['path']} (download once with cleanfid on a node with network)")
    elif exp_inc.get("sha256") and exp_inc["sha256"] != got["inception"]["sha256"]:
        probs.append(f"Inception weights sha256 {got['inception']['sha256'][:12]} != pinned {exp_inc['sha256'][:12]}")
    if probs:
        raise RuntimeError("FID pipeline pin violated: " + "; ".join(probs))
    return got


def build_extractor(device):
    """clean-fid 'clean' extractor: TorchScript TF-Inception, weights from INCEPTION_DIR, no internal resize."""
    from cleanfid.inception_torchscript import InceptionV3W
    INCEPTION_DIR.mkdir(parents=True, exist_ok=True)
    model = InceptionV3W(str(INCEPTION_DIR), download=not (INCEPTION_DIR / INCEPTION_FILE).is_file(),
                         resize_inside=False).to(device).eval()
    return model


def inception_features(imgs_u8: np.ndarray, extractor, device, batch: int = 200) -> np.ndarray:
    """uint8 (N, H, W, 3) -> float32 (N, 2048): a single clean resize per image, then the pinned network."""
    import torch
    from cleanfid.resize import build_resizer
    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}")
    resizer = build_resizer("clean")
    out = []
    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)]]           # float32 (299, 299, 3) in [0, 255]
            x = torch.from_numpy(np.stack(xs).transpose(0, 3, 1, 2).copy()).to(device)
            if tuple(x.shape[-2:]) != (IMAGE_SIZE, IMAGE_SIZE):
                raise AssertionError("resizer did not produce 299x299")
            out.append(extractor(x).cpu().numpy())
    return np.concatenate(out).astype(np.float32)


def fid_from_features(f1: np.ndarray, f2: np.ndarray) -> float:
    from cleanfid.fid import frechet_distance
    f1 = np.asarray(f1, np.float64); f2 = np.asarray(f2, np.float64)
    return float(frechet_distance(f1.mean(0), np.cov(f1, rowvar=False), f2.mean(0), np.cov(f2, rowvar=False)))


def kid_from_features(f1: np.ndarray, f2: np.ndarray, num_subsets: int = 100, max_subset_size: int = 1000) -> float:
    from cleanfid.fid import kernel_distance
    return float(kernel_distance(np.asarray(f1, np.float64), np.asarray(f2, np.float64),
                                 num_subsets=int(num_subsets), max_subset_size=int(max_subset_size)))


# ── csFID feature extraction ────────────────────────────────────────────────────────────────────────────────

def _pre_mixed6_stages(layers):
    """The layer sequence of layers.forward up to (not including) mixed_6 -- mirrors its TorchScript source."""
    return [layers.conv, layers.conv_1, layers.conv_2, layers.pool0, layers.conv_3, layers.conv_4, layers.pool1,
            layers.mixed, layers.mixed_1, layers.mixed_2, layers.mixed_3, layers.mixed_4, layers.mixed_5]


def _post_mixed6_stages(layers):
    return [layers.mixed_6, layers.mixed_7, layers.mixed_8, layers.mixed_9, layers.mixed_10, layers.pool2]


def spatial_tap(extractor, x_normalized):
    """x already resized to 299x299 and normalized as InceptionV3W does ((x - 128) / 128) -> the mixed_6/conv
    activation (N, 192, 17, 17)."""
    L = extractor.layers
    h = x_normalized
    for st in _pre_mixed6_stages(L):
        h = st.forward(h)
    return L.mixed_6.conv.forward(h)


def spatial_features(imgs_u8: np.ndarray, extractor, device, batch: int = 200) -> np.ndarray:
    """uint8 (N, H, W, 3) -> float32 (N, 2023): the canonical sFID feature. Same clean resize and normalization
    as inception_features; only the tapped tensor differs."""
    import torch
    from cleanfid.resize import build_resizer
    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}")
    resizer = build_resizer("clean")
    out = []
    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:]) != (IMAGE_SIZE, IMAGE_SIZE):
                raise AssertionError("resizer did not produce 299x299")
            sp = spatial_tap(extractor, (x - 128) / 128)[:, :SPATIAL_CHANNELS]
            out.append(sp.permute(0, 2, 3, 1).reshape(sp.shape[0], -1).cpu().numpy())
    feats = np.concatenate(out).astype(np.float32)
    if feats.shape[1] != SPATIAL_DIM:
        raise AssertionError(f"spatial feature dim {feats.shape[1]} != {SPATIAL_DIM}")
    return feats


def verify_spatial_tap(extractor=None) -> dict:
    """Check on CPU that (a) the stepwise traversal reproduces layers.forward exactly (so the tap sits on the
    cFID path) and (b) the tapped tensor is the canonical 17x17 mixed_6/conv representation. Raises on failure;
    returns the record that MANIFEST_CSFID.json carries."""
    import torch
    ext = extractor if extractor is not None else build_extractor(torch.device("cpu"))
    L = ext.layers
    if not (hasattr(L, "mixed_5") and hasattr(L, "mixed_6") and hasattr(L.mixed_6, "conv")):
        raise RuntimeError("this Inception build does not expose layers.mixed_6.conv -- canonical sFID unavailable")
    x = (torch.rand(3, 3, IMAGE_SIZE, IMAGE_SIZE) * 255.0 - 128) / 128
    with torch.no_grad():
        h = x
        for st in _pre_mixed6_stages(L):
            h = st.forward(h)
        mixed5 = h
        sp = L.mixed_6.conv.forward(mixed5)
        for st in _post_mixed6_stages(L):
            h = st.forward(h)
        diff = float((h.view(3, POOL3_DIM) - L.forward(x).view(3, POOL3_DIM)).abs().max())
    flat_dim = int(sp[:, :SPATIAL_CHANNELS].permute(0, 2, 3, 1).reshape(3, -1).shape[1])
    rec = {"mixed_5_output_shape_CHW": list(sp.shape[1:]) and list(mixed5.shape[1:]),
           "mixed_6_conv_output_shape_CHW": list(sp.shape[1:]), "channels_total": int(sp.shape[1]),
           "channels_kept": SPATIAL_CHANNELS, "feature_dim": flat_dim, "flatten": "NHWC",
           "stepwise_vs_layers_forward_max_abs_diff_cpu": diff, "bitexact_cpu": diff == 0.0}
    probs = []
    if tuple(mixed5.shape[1:]) != _MIXED5_SHAPE:
        probs.append(f"mixed_5 output {tuple(mixed5.shape[1:])} != canonical {_MIXED5_SHAPE}")
    if tuple(sp.shape[2:]) != (SPATIAL_GRID, SPATIAL_GRID):
        probs.append(f"mixed_6/conv spatial {tuple(sp.shape[2:])} != {SPATIAL_GRID}x{SPATIAL_GRID}")
    if flat_dim != SPATIAL_DIM:
        probs.append(f"feature dim {flat_dim} != {SPATIAL_DIM}")
    if diff != 0.0:
        probs.append(f"stepwise traversal differs from layers.forward on CPU by {diff}")
    if probs:
        raise RuntimeError("spatial tap verification failed: " + "; ".join(probs))
    return rec