"""
Produce stratified train/val/test splits from the master manifest.

Stratification preserves the per-class ratio across splits — important when
the classes are imbalanced (and they always are, eventually).

Determinism: a fixed seed plus sha256 sorting means re-running on the same
manifest produces identical splits. This matters for reproducibility.

Usage
-----
    python scripts/dataset/split.py \
        --manifest data/manifest.csv \
        --out-dir data \
        --val 0.1 --test 0.1

Generator-aware Stage 3A split:

    python scripts/dataset/split.py \
        --manifest data/manifest.csv \
        --out-dir data \
        --stratify-by class-generator

Held-out generator evaluation split:

    python scripts/dataset/split.py \
        --manifest data/manifest.csv \
        --out-dir data \
        --stratify-by class-generator \
        --holdout-generator sdxl

Content-disjoint held-out split (recommended for generalisation claims):

    python scripts/dataset/split.py \
        --manifest data/manifest.csv \
        --out-dir data/splits_v2 \
        --stratify-by class-generator \
        --holdout-generator sdxl \
        --group-by content-id

Why --group-by content-id exists
--------------------------------
Output filenames are content-addressed from the (prompt, seed) pair, and every
Stage 3A generator consumed the same `prompts.txt` with the same deterministic
seed derivation. So `<stem>.png` denotes the same prompt and seed in the SDXL,
SD 3.5 and AuraFlow sets.

`--holdout-generator` alone takes *every* row of that generator, which means
the held-out set can depict prompts that a sibling generator already
contributed to training. That is exactly what happened in Stage 3A: all 20,000
held-out SDXL images shared their (prompt, seed) with an SD 3.5 image in
train/val/test, so the split measured "unseen generator, **seen** content"
while being described as a true generalisation test.

With `--group-by content-id`, a (prompt, seed) identity is assigned to exactly
one partition, and two held-out files are written:

  heldout.csv               unseen generator AND unseen content (honest test)
  heldout_seen_content.csv  unseen generator, seen content (legacy-comparable)

Rows of non-held-out generators that fall in the reserved held-out content
partition are dropped from train/val/test and reported, since including them
would reintroduce the leak.

Default remains `--group-by none`, which reproduces the legacy behaviour
byte-for-byte.
"""
from __future__ import annotations

import argparse
import csv
import random
from collections import defaultdict
from pathlib import Path


def _content_id(row: dict[str, str]) -> str:
    """Identity of the (prompt, seed) pair behind a content-addressed filename.

    The same stem denotes the same prompt and seed across every Stage 3A
    generator, so this groups sibling renderings of identical content.
    """
    return Path(row["path"]).stem


def _generator_label(row: dict[str, str]) -> str:
    """Return the best available generator label for an AI row."""
    return row.get("generator") or row.get("source") or "unknown"


def _split_key(row: dict[str, str], stratify_by: str) -> str:
    """Return the grouping key used for stratified splitting."""
    cls = row["class"]
    if stratify_by == "class-generator" and cls == "ai_generated":
        return f"{cls}:{_generator_label(row)}"
    return cls


def _write_csv(path: Path, fieldnames: list[str], rows: list[dict[str, str]]) -> None:
    with path.open("w", newline="") as fh:
        writer = csv.DictWriter(fh, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)


def _split_content_disjoint(
    args: argparse.Namespace,
    rows: list[dict[str, str]],
    fieldnames: list[str],
    holdout_generators: set[str],
) -> None:
    """Partition by (prompt, seed) identity so held-out content is truly unseen."""
    ai_rows = [r for r in rows if r["class"] == "ai_generated"]
    real_rows = [r for r in rows if r["class"] != "ai_generated"]

    groups: dict[str, list[dict[str, str]]] = defaultdict(list)
    for r in ai_rows:
        groups[_content_id(r)].append(r)

    group_ids = sorted(groups)
    rng = random.Random(args.seed)
    rng.shuffle(group_ids)

    n_groups = len(group_ids)
    n_holdout = int(round(n_groups * args.holdout_frac)) if holdout_generators else 0
    holdout_ids = set(group_ids[:n_holdout])
    remaining = group_ids[n_holdout:]

    n_test = int(round(len(remaining) * args.test))
    n_val = int(round(len(remaining) * args.val))
    partition = {gid: "test" for gid in remaining[:n_test]}
    partition.update({gid: "val" for gid in remaining[n_test : n_test + n_val]})
    partition.update({gid: "train" for gid in remaining[n_test + n_val :]})

    buckets: dict[str, list[dict[str, str]]] = defaultdict(list)
    dropped: dict[str, int] = defaultdict(int)
    for gid, items in groups.items():
        in_holdout_content = gid in holdout_ids
        for r in items:
            is_holdout_gen = _generator_label(r) in holdout_generators
            if is_holdout_gen:
                buckets["heldout" if in_holdout_content else "heldout_seen_content"].append(r)
            elif in_holdout_content:
                # Including this would put the held-out prompt back into
                # training via a sibling generator — the exact Stage 3A leak.
                dropped[_generator_label(r)] += 1
            else:
                buckets[partition[gid]].append(r)

    # Real photos share no content identity with generated images, so they are
    # split by the ordinary stratified fractions.
    real_sorted = sorted(real_rows, key=lambda r: r["sha256"])
    rng_real = random.Random(args.seed)
    rng_real.shuffle(real_sorted)
    n = len(real_sorted)
    n_test_r = int(round(n * args.test))
    n_val_r = int(round(n * args.val))
    buckets["test"].extend(real_sorted[:n_test_r])
    buckets["val"].extend(real_sorted[n_test_r : n_test_r + n_val_r])
    buckets["train"].extend(real_sorted[n_test_r + n_val_r :])

    print(f"content-id groups: {n_groups}")
    print(f"  reserved for held-out content: {len(holdout_ids)}")
    print(f"  train/val/test groups: {len(remaining) - n_test - n_val}/{n_val}/{n_test}")
    if dropped:
        total_dropped = sum(dropped.values())
        print(
            f"  dropped {total_dropped} non-held-out rows whose content is "
            "reserved (prevents leakage):"
        )
        for gen, count in sorted(dropped.items()):
            print(f"    {gen}: {count}")

    args.out_dir.mkdir(parents=True, exist_ok=True)
    for name in ("train", "val", "test", "heldout", "heldout_seen_content"):
        rs = buckets.get(name, [])
        if not rs:
            continue
        rs.sort(key=lambda r: (r["sha256"], r["path"]))
        path = args.out_dir / f"{name}.csv"
        _write_csv(path, fieldnames, rs)
        counts: dict[str, int] = defaultdict(int)
        for r in rs:
            key = (
                _generator_label(r) if r["class"] == "ai_generated" else "authentic"
            )
            counts[key] += 1
        breakdown = "  ".join(f"{k}={v}" for k, v in sorted(counts.items()))
        print(f"  wrote {path} ({len(rs)} rows)  {breakdown}")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument("--out-dir", type=Path, required=True)
    parser.add_argument("--val", type=float, default=0.1)
    parser.add_argument("--test", type=float, default=0.1)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument(
        "--stratify-by",
        choices=["class", "class-generator"],
        default="class",
        help=(
            "Stratification mode. 'class' preserves legacy behavior; "
            "'class-generator' additionally balances AI generators."
        ),
    )
    parser.add_argument(
        "--holdout-generator",
        action="append",
        default=[],
        help=(
            "AI generator to remove from train/val/test and write to "
            "heldout.csv. Can be passed multiple times."
        ),
    )
    parser.add_argument(
        "--group-by",
        choices=["none", "content-id"],
        default="none",
        help=(
            "'none' (default) reproduces legacy behaviour. 'content-id' keeps "
            "every generator's rendering of the same (prompt, seed) in one "
            "partition, so the held-out split is content-disjoint as well as "
            "generator-disjoint."
        ),
    )
    parser.add_argument(
        "--holdout-frac",
        type=float,
        default=0.1,
        help=(
            "Fraction of content-id groups reserved for the held-out content "
            "partition. Only used with --group-by content-id."
        ),
    )
    args = parser.parse_args()

    if args.val + args.test >= 1.0:
        raise ValueError("val + test must be < 1.0")

    with args.manifest.open() as fh:
        reader = csv.DictReader(fh)
        rows = list(reader)
        fieldnames = reader.fieldnames or []

    holdout_generators = set(args.holdout_generator)

    if args.group_by == "content-id":
        return _split_content_disjoint(args, rows, fieldnames, holdout_generators)

    split_candidates: list[dict[str, str]] = []
    holdout_rows: list[dict[str, str]] = []
    for r in rows:
        if r["class"] == "ai_generated" and _generator_label(r) in holdout_generators:
            holdout_rows.append(r)
        else:
            split_candidates.append(r)

    if holdout_generators and not holdout_rows:
        raise ValueError(
            "No rows matched --holdout-generator values: "
            f"{sorted(holdout_generators)}"
        )

    by_group: dict[str, list[dict[str, str]]] = defaultdict(list)
    for r in split_candidates:
        by_group[_split_key(r, args.stratify_by)].append(r)

    rng = random.Random(args.seed)
    train_rows, val_rows, test_rows = [], [], []
    for group, items in sorted(by_group.items()):
        # Sort by sha256 for determinism, then shuffle with seeded RNG.
        items.sort(key=lambda r: r["sha256"])
        rng.shuffle(items)
        n = len(items)
        n_test = int(round(n * args.test))
        n_val = int(round(n * args.val))
        test_rows.extend(items[:n_test])
        val_rows.extend(items[n_test : n_test + n_val])
        train_rows.extend(items[n_test + n_val :])
        print(
            f"  {group}: total={n}  "
            f"train={n - n_test - n_val}  val={n_val}  test={n_test}"
        )

    args.out_dir.mkdir(parents=True, exist_ok=True)
    for name, rs in [("train", train_rows), ("val", val_rows), ("test", test_rows)]:
        path = args.out_dir / f"{name}.csv"
        _write_csv(path, fieldnames, rs)
        print(f"  wrote {path} ({len(rs)} rows)")

    if holdout_generators:
        path = args.out_dir / "heldout.csv"
        holdout_rows.sort(key=lambda r: (r["sha256"], r["path"]))
        _write_csv(path, fieldnames, holdout_rows)
        print(
            f"  wrote {path} ({len(holdout_rows)} rows) "
            f"for held-out generators: {sorted(holdout_generators)}"
        )


if __name__ == "__main__":
    main()