from __future__ import annotations

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


def classify_transformation(pair: dict) -> str:
    age_changed = pair.get("source_age") != pair.get("target_age")
    sex_changed = pair.get("source_sex") != pair.get("target_sex")
    if age_changed and sex_changed:
        return "intersectional"
    if age_changed:
        return "age_only"
    if sex_changed:
        return "sex_only"
    return "none"


def stratified_sample(
    pairs: list[dict],
    n_per_stratum: int,
    seed: int,
) -> list[int]:
    rng = random.Random(seed)

    by_stratum: dict[str, list[int]] = defaultdict(list)
    for i, pair in enumerate(pairs):
        stratum = classify_transformation(pair)
        if stratum != "none":
            by_stratum[stratum].append(i)

    selected: list[int] = []
    for stratum in ["age_only", "sex_only", "intersectional"]:
        pool = by_stratum[stratum]
        rng.shuffle(pool)
        n = min(n_per_stratum, len(pool))
        selected.extend(pool[:n])
        if n < n_per_stratum:
            print(
                f"Warning: only {n} pairs available for {stratum} "
                f"(requested {n_per_stratum})"
            )

    return selected


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Stratified sampling of pairs for radiologist validation"
    )
    parser.add_argument("--pairs", type=str, required=True)
    parser.add_argument("--n-per-stratum", type=int, default=150)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--output", type=str, default="data/pairs.json")
    args = parser.parse_args()

    with open(args.pairs) as f:
        all_pairs = json.load(f)

    indices = stratified_sample(all_pairs, args.n_per_stratum, args.seed)
    sampled = [all_pairs[i] for i in indices]

    for pair in sampled:
        pair["_transform_type"] = classify_transformation(pair)

    counts: dict[str, int] = defaultdict(int)
    for pair in sampled:
        counts[pair["_transform_type"]] += 1

    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    with open(output_path, "w") as f:
        json.dump(sampled, f, indent=2)

    print(f"Sampled {len(sampled)} pairs from {len(all_pairs)} total")
    for stratum, count in sorted(counts.items()):
        print(f"  {stratum}: {count}")
    print(f"Saved to: {output_path}")


if __name__ == "__main__":
    main()