Blame for scripts/precompute_frequency_features.py

"""
Pre-compute Stage 3B frequency features for the dataset splits.

Mirrors `scripts/precompute_embeddings.py` in shape (same split CSVs, same
`--data-root` resolution, same label map, resumable shards) but emits the
15-dim `freq-v1` feature vector from
`deepfake_scanner.detectors.frequency_features` instead of CLIP embeddings.

Why a separate offline step?
----------------------------
Feature extraction is pure-numpy and CPU-cheap (~a few ms/image), but the
full multi-generator dataset is ~100k+ images, so caching the features once
lets the classifier (S3B-07) iterate in seconds without re-reading images.

Output format
-------------
One `{split}_freq.npz` per split in `--out-dir`:
  - features:        float32 (N, 16)   freq-v1 feature vectors
  - labels:          int8    (N,)      0 = authentic, 1 = ai_generated
  - paths:           unicode (N,)      manifest-relative image paths
  - sources:         unicode (N,)
  - generators:      unicode (N,)
  - model_families:  unicode (N,)
  - augmentations:   unicode (N,)
  - feature_schema:  scalar str        FEATURE_SCHEMA_VERSION ("freq-v1")
  - feature_names:   unicode (16,)     FEATURE_NAMES, for validation

Resumability
------------
Each split is encoded in shards (`{split}_freq_shards/shard_NNNNNNNN.npz`);
on restart we resume from the first missing contiguous shard. Final files
are concatenated at the end and the shard dir removed. Idempotent: an
existing final `.npz` is skipped unless `--force`.

Usage
-----
    python scripts/precompute_frequency_features.py \\
        --splits-dir data/splits \\
        --out-dir data/freq_features \\
        --data-root data

    # Include the augmented splits used for robustness training/eval:
    python scripts/precompute_frequency_features.py \\
        --splits-dir data/splits \\
        --out-dir data/freq_features \\
        --data-root data \\
        --splits train val test heldout test_augmented train_augmented
"""
from __future__ import annotations

import argparse
import csv
import shutil
import time
from dataclasses import dataclass
from pathlib import Path

import numpy as np
from PIL import Image

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

# Keep label mapping identical to precompute_embeddings.py / train_head.py:
# logits[..., 1] = ai_generated.
_LABEL_MAP = {"authentic": 0, "ai_generated": 1}

# Images per shard. Feature extraction is cheap, so use a larger shard than
# the embeddings script; a crash loses at most a few seconds of work.
SHARD_SIZE = 2000


@dataclass(frozen=True)
class FeatureItem:
    label: int
    manifest_path: str
    image_path: Path
    source: str
    generator: str
    model_family: str
    augmentation: str


def _read_split(csv_path: Path, data_root: Path) -> list[FeatureItem]:
    items: list[FeatureItem] = []
    with csv_path.open(encoding="utf-8") as fh:
        reader = csv.DictReader(fh)
        for row in reader:
            label = _LABEL_MAP[row["class"]]
            manifest_path = row["path"]
            items.append(
                FeatureItem(
                    label=label,
                    manifest_path=manifest_path,
                    image_path=data_root / manifest_path,
                    source=row.get("source", ""),
                    generator=row.get("generator", ""),
                    model_family=row.get("model_family", ""),
                    augmentation=row.get("augmentation", ""),
                )
            )
    return items


def _resume_cursor(shard_dir: Path) -> int:
    """Count the contiguous valid shard prefix; delete a trailing corrupt one."""
    if not shard_dir.exists():
        return 0
    cursor = 0
    for p in sorted(shard_dir.glob("shard_*.npz")):
        try:
            with np.load(p) as z:
                _ = z["features"], z["labels"]
            cursor += int(np.load(p)["features"].shape[0])
        except Exception:
            print(f"  removing corrupt shard {p.name}")
            p.unlink()
            break
    return cursor


def _save_shard_atomic(shard_path: Path, **arrays) -> None:
    tmp = shard_path.with_suffix(".tmp.npz")
    np.savez_compressed(tmp, **arrays)
    tmp.replace(shard_path)


def _concat_shards(shard_dir: Path, out_path: Path, expected_total: int) -> None:
    paths = sorted(shard_dir.glob("shard_*.npz"))
    keys = [
        "features",
        "labels",
        "paths",
        "sources",
        "generators",
        "model_families",
        "augmentations",
    ]
    acc: dict[str, list[np.ndarray]] = {k: [] for k in keys}
    for p in paths:
        with np.load(p) as z:
            for k in keys:
                acc[k].append(z[k].copy())
    merged = {k: np.concatenate(acc[k], axis=0) for k in keys}
    if len(merged["labels"]) != expected_total:
        raise RuntimeError(
            f"Concatenated {len(merged['labels'])} rows but expected "
            f"{expected_total}. Refusing to write {out_path}."
        )
    np.savez_compressed(
        out_path,
        feature_schema=np.asarray(FEATURE_SCHEMA_VERSION),
        feature_names=np.asarray(FEATURE_NAMES),
        **merged,
    )
    shutil.rmtree(shard_dir)


def _encode_split(
    split: str,
    splits_dir: Path,
    out_dir: Path,
    data_root: Path,
    force: bool,
) -> None:
    csv_path = splits_dir / f"{split}.csv"
    if not csv_path.exists():
        print(f"Skipping {split} — {csv_path} not found.")
        return

    out_path = out_dir / f"{split}_freq.npz"
    shard_dir = out_dir / f"{split}_freq_shards"

    if force:
        if out_path.exists():
            out_path.unlink()
        if shard_dir.exists():
            shutil.rmtree(shard_dir)

    if out_path.exists():
        print(f"Skipping {split} — {out_path} already exists. Use --force.")
        return

    items = _read_split(csv_path, data_root)
    n = len(items)
    print(f"{split}: {n} images")

    cursor = _resume_cursor(shard_dir)
    if cursor:
        print(f"  resuming from {cursor}/{n}")
    shard_dir.mkdir(parents=True, exist_ok=True)

    t0 = time.time()
    while cursor < n:
        shard_end = min(cursor + SHARD_SIZE, n)
        chunk = items[cursor:shard_end]
        feats = np.empty((len(chunk), N_FEATURES), dtype=np.float32)
        for i, item in enumerate(chunk):
            with Image.open(item.image_path) as img:
                feats[i] = extract_features(img)
        _save_shard_atomic(
            shard_dir / f"shard_{cursor:08d}.npz",
            features=feats,
            labels=np.asarray([it.label for it in chunk], dtype=np.int8),
            paths=np.asarray([it.manifest_path for it in chunk]),
            sources=np.asarray([it.source for it in chunk]),
            generators=np.asarray([it.generator for it in chunk]),
            model_families=np.asarray([it.model_family for it in chunk]),
            augmentations=np.asarray([it.augmentation for it in chunk]),
        )
        cursor = shard_end
        elapsed = time.time() - t0
        rate = cursor / elapsed if elapsed else 0.0
        eta = (n - cursor) / rate if rate else 0.0
        print(
            f"  {cursor}/{n}  ({rate:.0f} img/s, eta {eta / 60:.1f} min, "
            f"shard saved)"
        )

    print(f"  concatenating shards into {out_path}...")
    _concat_shards(shard_dir, out_path, expected_total=n)
    print(f"  wrote {out_path}  ({out_path.stat().st_size / 1e6:.1f} MB)")


def main() -> None:
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("--splits-dir", type=Path, required=True)
    parser.add_argument("--out-dir", type=Path, required=True)
    parser.add_argument("--data-root", type=Path, required=True)
    parser.add_argument(
        "--splits",
        nargs="+",
        default=["train", "val", "test", "heldout", "test_augmented"],
        help="Split names (without .csv) to encode.",
    )
    parser.add_argument(
        "--force",
        action="store_true",
        help="Recompute even if an output .npz already exists.",
    )
    args = parser.parse_args()

    args.out_dir.mkdir(parents=True, exist_ok=True)
    print(f"Feature schema: {FEATURE_SCHEMA_VERSION} ({N_FEATURES} features)")
    for split in args.splits:
        _encode_split(
            split=split,
            splits_dir=args.splits_dir,
            out_dir=args.out_dir,
            data_root=args.data_root,
            force=args.force,
        )


if __name__ == "__main__":
    main()