"""
QA report for cached Stage 3B frequency features (S3B-06).

Reads the `{split}_freq.npz` files written by
`scripts/precompute_frequency_features.py` and produces a small JSON+text
report so we can sanity-check the features BEFORE training a classifier:

  - per-feature mean/std, split by class (authentic vs ai_generated)
  - per-feature mean by generator (where present)
  - NaN/Inf counts (must be zero — the extractor forces finite, this verifies)
  - extreme-outlier counts (|z| > 8 against the split's own mean/std)
  - a crude per-feature separability score: |mean_ai - mean_real| / pooled_std,
    which flags which features carry class signal at all (a fast pre-check of
    whether the frequency approach has any chance before we train).

No new dependencies (numpy only). Read-only over the feature cache.

Usage
-----
    python scripts/report_frequency_features.py \\
        --feat-dir data/freq_features \\
        --splits train val test heldout test_augmented \\
        --report-out data/reports/freq_v1_feature_qa.json
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np

from deepfake_scanner.detectors.frequency_features import (
    FEATURE_NAMES,
    FEATURE_SCHEMA_VERSION,
)


def _load(feat_dir: Path, split: str):
    path = feat_dir / f"{split}_freq.npz"
    if not path.exists():
        return None
    with np.load(path, allow_pickle=False) as z:
        return {
            "features": z["features"].astype(np.float64),
            "labels": z["labels"].astype(np.int64),
            "generators": np.asarray([str(g) for g in z["generators"]]),
            "schema": str(z["feature_schema"]),
        }


def _separability(feats: np.ndarray, labels: np.ndarray) -> list[float]:
    """|mean_ai - mean_real| / pooled_std per feature. Higher = more signal."""
    real = feats[labels == 0]
    ai = feats[labels == 1]
    if len(real) == 0 or len(ai) == 0:
        return [0.0] * feats.shape[1]
    pooled = np.sqrt((real.var(axis=0) + ai.var(axis=0)) / 2.0) + 1e-8
    return (np.abs(ai.mean(axis=0) - real.mean(axis=0)) / pooled).tolist()


def main() -> None:
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("--feat-dir", type=Path, required=True)
    parser.add_argument(
        "--splits",
        nargs="+",
        default=["train", "val", "test", "heldout", "test_augmented"],
    )
    parser.add_argument("--report-out", type=Path, default=None)
    args = parser.parse_args()

    report: dict = {"feature_schema": FEATURE_SCHEMA_VERSION, "splits": {}}

    for split in args.splits:
        data = _load(args.feat_dir, split)
        if data is None:
            print(f"skip {split} (no feature file)")
            continue
        if data["schema"] != FEATURE_SCHEMA_VERSION:
            raise ValueError(
                f"{split}: schema {data['schema']!r} != {FEATURE_SCHEMA_VERSION!r}"
            )
        x = data["features"]
        y = data["labels"]
        n_nan = int(np.isnan(x).sum())
        n_inf = int(np.isinf(x).sum())
        mean = x.mean(axis=0)
        std = x.std(axis=0) + 1e-8
        z = np.abs((x - mean) / std)
        n_outliers = int((z > 8).sum())
        sep = _separability(x, y)

        split_rep = {
            "n": int(len(y)),
            "class_counts": {
                "authentic": int((y == 0).sum()),
                "ai_generated": int((y == 1).sum()),
            },
            "nan_count": n_nan,
            "inf_count": n_inf,
            "extreme_outlier_count": n_outliers,
            "per_feature": {
                FEATURE_NAMES[i]: {
                    "mean": float(mean[i]),
                    "std": float(std[i] - 1e-8),
                    "separability": float(sep[i]),
                }
                for i in range(len(FEATURE_NAMES))
            },
        }
        report["splits"][split] = split_rep

        # Console summary: top separating features on this split.
        order = np.argsort(sep)[::-1]
        top = ", ".join(f"{FEATURE_NAMES[i]}={sep[i]:.2f}" for i in order[:5])
        print(
            f"[{split}] n={len(y)} nan={n_nan} inf={n_inf} "
            f"outliers={n_outliers} | top-sep: {top}"
        )

    if args.report_out is not None:
        args.report_out.parent.mkdir(parents=True, exist_ok=True)
        with args.report_out.open("w", encoding="utf-8") as fh:
            json.dump(report, fh, indent=2, sort_keys=True)
        print(f"saved report to {args.report_out}")


if __name__ == "__main__":
    main()