"""Conditional image metrics per test condition: csFID (the paper's metric), cFID and cKID.

  generated_states(samples, variant, packaged, n)   RAW generated states of every condition from a samples directory
  truth_features / truth_spatial_features           cached features of the first n_truth bank members of one condition
  evaluate_csfid_run(...)                           csFID of all conditions -> csfid_per_ic.csv, csfid_summary.json
  evaluate_run(...)                                 cFID and cKID (pool3 features) -> per_ic_fid.csv, summary.json

The truth bank holds ``<bank_dir>/group_<gid:04d>.npz`` with ``vecs`` (n_members, n_snapshots, D) in raw units; the
first n_truth members are the reference ensemble of every method.
"""
from __future__ import annotations

import json
from datetime import datetime, timezone
from pathlib import Path

import numpy as np
import pandas as pd

from . import backend as BE
from .adapters import FIDAdapter


def generated_states(samples_dir: Path, variant: str, packaged: Path, n: int) -> dict:
    """{gid: RAW generated states (n, D)} of ``samples_dir/<variant>/group_<gid>.npz`` (``gen_xT_norm``, un-normalized
    with the package's normalization.npz)."""
    z = np.load(Path(packaged) / "normalization.npz")
    mean, std = np.asarray(z["mean"], np.float64), np.asarray(z["std"], np.float64)
    out = {}
    for p in sorted((Path(samples_dir) / variant).glob("group_[0-9][0-9][0-9][0-9].npz")):
        with np.load(p) as f:
            g = np.asarray(f["gen_xT_norm"], np.float64)[: int(n)]
            out[int(f["gid"])] = g * std + mean
    return out


def bank_members(bank_dir: Path, gid: int, t_idx: int, n: int | None = None) -> np.ndarray:
    """RAW states (n, D) of the first n members of test condition ``gid`` at snapshot ``t_idx``."""
    p = Path(bank_dir) / f"group_{int(gid):04d}.npz"
    if not p.is_file():
        raise FileNotFoundError(f"truth bank has no group_{int(gid):04d}.npz under {bank_dir}")
    with np.load(p) as z:
        vecs = z["vecs"]
        n_have = int(vecs.shape[0])
        if n is not None and n_have < int(n):
            raise ValueError(f"bank group {gid}: {n_have} members < required {n}")
        return np.asarray(vecs[: (int(n) if n is not None else n_have), int(t_idx), :], dtype=np.float64)


def bank_inventory(bank_dir: Path, gids: list[int]) -> dict:
    out = {}
    for g in gids:
        p = Path(bank_dir) / f"group_{int(g):04d}.npz"
        if not p.is_file():
            out[int(g)] = 0
            continue
        with np.load(p) as z:
            out[int(g)] = int(z["vecs"].shape[0])
    return out


def features_for(states_raw: np.ndarray, adapter: FIDAdapter, extractor, device, batch: int = 200) -> tuple[np.ndarray, float]:
    imgs, clipped = adapter.to_images(states_raw)
    return BE.inception_features(imgs, extractor, device, batch=batch), clipped


def truth_features(bank_dir: Path, gid: int, t_idx: int, n_truth: int, adapter: FIDAdapter, extractor, device,
                   cache_dir: Path, batch: int = 200) -> tuple[np.ndarray, dict]:
    """Cached per (gid, n_truth, adapter clip_c): the truth side does not depend on the model."""
    cache_dir = Path(cache_dir); cache_dir.mkdir(parents=True, exist_ok=True)
    tag = f"gid{int(gid):04d}_t{int(t_idx)}_n{int(n_truth)}_c{adapter.params['clip_c']:.6g}"
    p = cache_dir / f"truth_{tag}.npz"
    if p.is_file():
        with np.load(p) as z:
            return np.asarray(z["feats"]), {"cache": str(p), "clipped_frac": float(z["clipped_frac"]), "reused": True}
    states = bank_members(bank_dir, gid, t_idx, n_truth)
    feats, clipped = features_for(states, adapter, extractor, device, batch)
    np.savez(p, feats=feats, clipped_frac=np.float64(clipped), gid=np.int64(gid), n=np.int64(n_truth))
    return feats, {"cache": str(p), "clipped_frac": clipped, "reused": False}


def spatial_features_for(states_raw: np.ndarray, adapter: FIDAdapter, extractor, device, batch: int = 200) -> tuple[np.ndarray, float]:
    imgs, clipped = adapter.to_images(states_raw)
    return BE.spatial_features(imgs, extractor, device, batch=batch), clipped


def truth_spatial_features(bank_dir: Path, gid: int, t_idx: int, n_truth: int, adapter: FIDAdapter, extractor, device,
                           cache_dir: Path, batch: int = 200) -> tuple[np.ndarray, dict]:
    """Like truth_features, for the spatial feature (a separate cache file)."""
    cache_dir = Path(cache_dir); cache_dir.mkdir(parents=True, exist_ok=True)
    tag = f"gid{int(gid):04d}_t{int(t_idx)}_n{int(n_truth)}_c{adapter.params['clip_c']:.6g}"
    p = cache_dir / f"truth_spatial_{tag}.npz"
    if p.is_file():
        with np.load(p) as z:
            return np.asarray(z["feats"]), {"cache": str(p), "clipped_frac": float(z["clipped_frac"]), "reused": True}
    states = bank_members(bank_dir, gid, t_idx, n_truth)
    feats, clipped = spatial_features_for(states, adapter, extractor, device, batch)
    np.savez(p, feats=feats, clipped_frac=np.float64(clipped), gid=np.int64(gid), n=np.int64(n_truth))
    return feats, {"cache": str(p), "clipped_frac": clipped, "reused": False}


def fid_kid(f_truth: np.ndarray, f_gen: np.ndarray, kid_cfg: dict) -> dict:
    return {"fid": BE.fid_from_features(f_truth, f_gen),
            "kid": BE.kid_from_features(f_truth, f_gen, int(kid_cfg.get("num_subsets", 100)),
                                        int(kid_cfg.get("max_subset_size", 1000)))}


def summarize_per_ic(df: pd.DataFrame, cols=("fid", "kid")) -> dict:
    out = {"n_ic": int(len(df))}
    for c in cols:
        v = df[c].to_numpy(dtype=float); v = v[np.isfinite(v)]
        out[c] = {"mean": float(v.mean()) if len(v) else float("nan"),
                  "sem": float(v.std(ddof=1) / np.sqrt(len(v))) if len(v) > 1 else float("nan"),
                  "sd": float(v.std(ddof=1)) if len(v) > 1 else float("nan"), "n": int(len(v))}
    return out


def evaluate_run(gen_by_gid: dict, bank_dir: Path, t_idx: int, n_truth: int, adapter: FIDAdapter, out_dir: Path,
                 cache_dir: Path, fid_cfg: dict, device, provenance: dict | None = None) -> dict:
    """cFID and cKID. ``gen_by_gid``: {gid: RAW generated states (n_gen, D)}; writes out_dir/{per_ic_fid.csv,
    summary.json, MANIFEST_FID.json, features/gen_gid<gid>.npz} and returns the summary (mean ± SEM over conditions)."""
    out_dir = Path(out_dir); (out_dir / "features").mkdir(parents=True, exist_ok=True)
    pins = BE.check_pins(fid_cfg)
    extractor = BE.build_extractor(device)
    batch = int(fid_cfg.get("feature_batch", 200))
    rows, prov_ics = [], {}
    for gid in sorted(gen_by_gid):
        gen = np.asarray(gen_by_gid[gid], dtype=np.float64)
        finite = np.isfinite(gen).all(axis=1)
        n_bad = int((~finite).sum())
        ft, tinfo = truth_features(bank_dir, gid, t_idx, n_truth, adapter, extractor, device, cache_dir, batch)
        if n_bad == len(gen):
            rows.append({"gid": int(gid), "fid": float("nan"), "kid": float("nan"), "n_truth": int(len(ft)),
                         "n_gen": 0, "n_gen_nonfinite": n_bad, "clipped_frac_gen": float("nan"), "clipped_frac_truth": tinfo["clipped_frac"]})
            continue
        fg, clipped_g = features_for(gen[finite], adapter, extractor, device, batch)
        np.savez(out_dir / "features" / f"gen_gid{int(gid):04d}.npz", feats=fg, gid=np.int64(gid))
        m = fid_kid(ft, fg, fid_cfg.get("kid", {}))
        rows.append({"gid": int(gid), **m, "n_truth": int(len(ft)), "n_gen": int(len(fg)), "n_gen_nonfinite": n_bad,
                     "clipped_frac_gen": clipped_g, "clipped_frac_truth": tinfo["clipped_frac"]})
        prov_ics[int(gid)] = tinfo
    df = pd.DataFrame(rows).sort_values("gid")
    df.to_csv(out_dir / "per_ic_fid.csv", index=False)
    summary = summarize_per_ic(df)
    summary["cFID"] = summary["fid"]; summary["cKID"] = summary["kid"]
    summary["n_truth_per_ic"] = int(n_truth)
    (out_dir / "summary.json").write_text(json.dumps(summary, indent=1) + "\n")
    man = {"schema": "fid_v1", "pins": pins, "adapter": adapter.params, "fid_settings": fid_cfg, "bank_dir": str(bank_dir),
           "t_idx": int(t_idx), "n_truth": int(n_truth), "truth_features": prov_ics, "provenance": provenance or {},
           "at": datetime.now(timezone.utc).isoformat(timespec="seconds")}
    (out_dir / "MANIFEST_FID.json").write_text(json.dumps(man, indent=1, default=str) + "\n")
    return summary


def evaluate_csfid_run(gen_by_gid: dict, bank_dir: Path, t_idx: int, n_truth: int, adapter: FIDAdapter, out_dir: Path,
                       cache_dir: Path, fid_cfg: dict, device, provenance: dict | None = None) -> dict:
    """csFID: the Frechet distance of the spatial Inception feature (Nash et al., 2021) between the generated states
    and the truth bank, per condition; writes out_dir/{csfid_per_ic.csv, csfid_summary.json, MANIFEST_CSFID.json} and
    returns the summary (mean ± SEM over conditions)."""
    out_dir = Path(out_dir); out_dir.mkdir(parents=True, exist_ok=True)
    pins = BE.check_pins(fid_cfg)
    sp = fid_cfg.get("spatial") or {}
    want = {"layer": BE.SPATIAL_LAYER, "channels": BE.SPATIAL_CHANNELS, "grid": BE.SPATIAL_GRID, "dim": BE.SPATIAL_DIM,
            "flatten": "NHWC", "pooling": "none"}
    bad = {k: (sp[k], v) for k, v in want.items() if k in sp and sp[k] != v}
    if bad:
        raise RuntimeError(f"the fid.spatial settings disagree with the csFID implementation: {bad}")
    pins["spatial_tap_check"] = BE.verify_spatial_tap()
    extractor = BE.build_extractor(device)
    batch = int(fid_cfg.get("feature_batch", 200))
    rows, prov_ics = [], {}
    for gid in sorted(gen_by_gid):
        gen = np.asarray(gen_by_gid[gid], dtype=np.float64)
        finite = np.isfinite(gen).all(axis=1)
        n_bad = int((~finite).sum())
        ft, tinfo = truth_spatial_features(bank_dir, gid, t_idx, n_truth, adapter, extractor, device, cache_dir, batch)
        if n_bad == len(gen):
            rows.append({"gid": int(gid), "sfid": float("nan"), "n_truth": int(len(ft)), "n_gen": 0, "n_gen_nonfinite": n_bad})
            continue
        fg, _ = spatial_features_for(gen[finite], adapter, extractor, device, batch)
        rows.append({"gid": int(gid), "sfid": BE.fid_from_features(ft, fg), "n_truth": int(len(ft)), "n_gen": int(len(fg)),
                     "n_gen_nonfinite": n_bad})
        prov_ics[int(gid)] = tinfo
    df = pd.DataFrame(rows).sort_values("gid")
    df.to_csv(out_dir / "csfid_per_ic.csv", index=False)
    summary = summarize_per_ic(df, cols=("sfid",))
    summary["csFID"] = summary["sfid"]
    summary["n_truth_per_ic"] = int(n_truth)
    (out_dir / "csfid_summary.json").write_text(json.dumps(summary, indent=1) + "\n")
    man = {"schema": "csfid_v1", "pins": pins, "adapter": adapter.params, "fid_settings": fid_cfg,
           "bank_dir": str(bank_dir), "t_idx": int(t_idx), "n_truth": int(n_truth), "truth_features": prov_ics,
           "provenance": provenance or {}, "at": datetime.now(timezone.utc).isoformat(timespec="seconds")}
    (out_dir / "MANIFEST_CSFID.json").write_text(json.dumps(man, indent=1, default=str) + "\n")
    return summary


def training_states_at(packaged_dir: Path, n_ic: int, members: int, snapshots) -> np.ndarray:
    """RAW states of the training pool (the first ``n_ic`` training conditions x the first ``members`` realizations) at
    the requested snapshots ('all' or a list of indices): the only data the adapter constants are fitted on."""
    from ..data.fine import FINE_FILE, _npz_member_memmap
    v = _npz_member_memmap(str(Path(packaged_dir) / FINE_FILE), "vecs_train")
    if v.shape[0] < n_ic or v.shape[1] < members:
        raise ValueError(f"vecs_train {tuple(v.shape)} smaller than the requested {n_ic} x {members} pool")
    if snapshots == "all":
        block = np.asarray(v[:n_ic, :members, :, :], dtype=np.float64)
    else:
        idx = [int(t) for t in (snapshots if isinstance(snapshots, (list, tuple)) else [snapshots])]
        block = np.asarray(v[:n_ic, :members, :, :][:, :, idx, :], dtype=np.float64)
    return block.reshape(-1, v.shape[-1])


def fit_adapter_from_training(adapter: FIDAdapter, packaged_dir: Path, n_ic: int, pool: dict, t_idx_final: int) -> dict:
    """``pool`` (dataset card): {members: int | 'all', snapshots: 'all' | 'final' | [indices], pixel_stride: int}."""
    from ..data.fine import FINE_FILE, _npz_member_memmap
    v = _npz_member_memmap(str(Path(packaged_dir) / FINE_FILE), "vecs_train")
    members = int(v.shape[1]) if str(pool.get("members", "all")) == "all" else int(pool["members"])
    snaps = pool.get("snapshots", "all")
    snaps = [int(t_idx_final)] if snaps == "final" else snaps
    states = training_states_at(packaged_dir, n_ic, members, snaps)
    return adapter.fit(states, pixel_stride=int(pool.get("pixel_stride", 1)),
                       source={"packaged_dataset_path": str(packaged_dir), "block": "vecs_train", "ics": f"canonical rows 0..{n_ic - 1}",
                               "members": f"j < {members}", "snapshots": ("all" if snaps == "all" else list(snaps)),
                               "pixel_stride": int(pool.get("pixel_stride", 1)), "n_states": int(len(states)),
                               "rule": "training data only; validation, test and generated data never enter"})