"""
Randomly sample N real chest X-rays by base ID, then build combined
train sets with each synthesis model.

Rules
-----
- A "base ID" is the trailing folder name in `file_name`, with any
  ``-1`` / ``-2`` / ``-N`` suffix stripped.  e.g. ``36255-1`` -> ``36255``.
- In the REAL jsonl we ONLY keep the base row (whose raw ID has no
  ``-N`` suffix).  Variant rows like ``36255-1`` are discarded entirely.
- We sample base IDs from the REAL jsonl, but only from base IDs that
  ALSO exist in every syn jsonl (cheff / roentgen / roentgenv2 / ours).
  This guarantees the four ``+model`` output files share the same ID set.
- Each syn file is assumed to have exactly one row per base ID.
- Output row counts:
    train_real{N}.json            : exactly N rows
    train_real{N}+<model>.json    : N real + N syn = 2N rows
- Outputs (jsonl, one record per line — keep extension as requested):
    train_real{N}.json                 # real rows only
    train_real{N}+cheff.json           # real rows + cheff syn rows
    train_real{N}+roentgen.json
    train_real{N}+roentgenv2.json
    train_real{N}+ours.json

Usage
-----
python sample_real500.py \
    --real      real.jsonl \
    --cheff     cheff.jsonl \
    --roentgen  roentgen.jsonl \
    --roentgenv2 roentgenv2.jsonl \
    --ours      ours.jsonl \
    --n 500 \
    --seed 42 \
    --outdir out/
"""
import argparse
import json
import os
import random
import re
import sys
from collections import defaultdict


# Strip trailing "-<digits>" from the base name (e.g. 36255-1 -> 36255).
# We deliberately only strip a SINGLE numeric suffix, since the variant
# scheme is -1, -2, ...  Ids that legitimately contain a dash followed by
# non-digits are left alone.
_SUFFIX_RE = re.compile(r"-\d+$")


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--real", required=True,
                   help="Path to the real jsonl.")
    p.add_argument("--cheff", required=True)
    p.add_argument("--roentgen", required=True)
    p.add_argument("--roentgenv2", required=True)
    p.add_argument("--ours", required=True)
    p.add_argument("--n", type=int, default=500,
                   help="Number of base IDs to sample.")
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--outdir", default=".")
    p.add_argument("--prefix", default=None,
                   help="Output filename prefix. Default: train_real{N}")
    return p.parse_args()


def extract_id(file_name: str) -> str:
    """
    Return the raw ID (folder name) from a record's ``file_name``.

    Examples
    --------
    >>> extract_id("AURAD_dataset/chestdet/dataset/train/36255-1/36255-1.png")
    '36255-1'
    >>> extract_id("AURAD_infer/det-train-real-mask-image-total-filter/AURA/dataset/train/46038/syn_46038.png")
    '46038'
    """
    # Take the directory of the image -> last path component.
    # This is robust to whether the basename has a "syn_" prefix or not.
    parent = os.path.basename(os.path.dirname(file_name))
    return parent


def base_id(raw_id: str) -> str:
    """Strip trailing ``-N`` numeric variant suffix."""
    return _SUFFIX_RE.sub("", raw_id)


def load_jsonl(path: str):
    """Yield (lineno, record_dict) for each non-empty line."""
    with open(path, "r", encoding="utf-8") as f:
        for i, line in enumerate(f, 1):
            line = line.strip()
            if not line:
                continue
            try:
                yield i, json.loads(line)
            except json.JSONDecodeError as e:
                print(f"[warn] {path}:{i} JSON decode error: {e}",
                      file=sys.stderr)


def index_real(path: str):
    """
    Return {base_id: record} where ``record`` is the base (non-variant) row,
    i.e. the row whose raw ID has no ``-N`` suffix.

    -1/-2/... variant rows are skipped entirely. If a base ID appears in
    the file only as variants (no base row), it is dropped and a warning
    is printed.
    """
    out = {}
    dropped_variant = 0
    seen_only_variant = defaultdict(int)
    for lineno, rec in load_jsonl(path):
        rid = extract_id(rec["file_name"])
        bid = base_id(rid)
        if rid != bid:
            # This is a -1 / -2 / ... variant; skip per spec.
            dropped_variant += 1
            seen_only_variant[bid] += 1
            continue
        if bid in out:
            # Duplicate base row for the same ID — keep first, warn.
            print(f"[warn] real:{lineno} duplicate base row for ID {bid}, "
                  f"keeping first.", file=sys.stderr)
            continue
        out[bid] = rec
        seen_only_variant.pop(bid, None)  # this id has a base row now

    if dropped_variant:
        print(f"[info] real: skipped {dropped_variant} -N variant row(s).",
              file=sys.stderr)
    if seen_only_variant:
        print(f"[warn] real: {len(seen_only_variant)} base ID(s) had only "
              f"variant rows and no base row — they are unusable and dropped.",
              file=sys.stderr)
    return out


def index_syn(path: str, label: str):
    """
    Return {base_id: record}.  Warn on duplicates and keep the first.
    """
    out = {}
    dup = 0
    for lineno, rec in load_jsonl(path):
        rid = extract_id(rec["file_name"])
        bid = base_id(rid)
        if bid in out:
            dup += 1
            continue
        out[bid] = rec
    if dup:
        print(f"[warn] {label}: {dup} duplicate base IDs skipped "
              f"(kept first occurrence).", file=sys.stderr)
    return out


def write_jsonl(path: str, records):
    with open(path, "w", encoding="utf-8") as f:
        for r in records:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")


def main():
    args = parse_args()
    os.makedirs(args.outdir, exist_ok=True)

    print("Loading real ...")
    real_idx = index_real(args.real)
    print(f"  real: {len(real_idx)} usable base IDs (base rows only)")

    syn_paths = {
        "cheff":      args.cheff,
        "roentgen":   args.roentgen,
        "roentgenv2": args.roentgenv2,
        "ours":       args.ours,
    }
    syn_idx = {}
    for name, p in syn_paths.items():
        print(f"Loading {name} ...")
        syn_idx[name] = index_syn(p, name)
        print(f"  {name}: {len(syn_idx[name])} unique base IDs")

    # Candidates = real ∩ every syn
    common = set(real_idx.keys())
    for name, idx in syn_idx.items():
        before = len(common)
        common &= set(idx.keys())
        print(f"After intersect with {name}: {len(common)} "
              f"(dropped {before - len(common)})")

    if len(common) < args.n:
        raise SystemExit(
            f"Only {len(common)} base IDs are common to all sources; "
            f"cannot sample {args.n}. Lower --n or check your inputs."
        )

    rng = random.Random(args.seed)
    # Sort for determinism before sampling (set order is nondeterministic).
    pool = sorted(common)
    sampled = rng.sample(pool, args.n)
    sampled_set = set(sampled)
    print(f"Sampled {len(sampled)} base IDs (seed={args.seed}).")

    # Build real records: exactly one row per sampled base ID (the base row).
    real_records = [real_idx[bid] for bid in sampled]
    print(f"Real rows in output: {len(real_records)} "
          f"(1 per base ID; -N variants discarded)")

    prefix = args.prefix or f"train_real{args.n}"

    # 1) real-only
    real_only_path = os.path.join(args.outdir, f"{prefix}.json")
    write_jsonl(real_only_path, real_records)
    print(f"Wrote {real_only_path}  ({len(real_records)} rows)")

    # 2) real + each syn
    for name in syn_paths:
        syn_records = [syn_idx[name][bid] for bid in sampled]
        combined = real_records + syn_records
        out_path = os.path.join(args.outdir, f"{prefix}+{name}.json")
        write_jsonl(out_path, combined)
        print(f"Wrote {out_path}  ({len(combined)} rows: "
              f"{len(real_records)} real + {len(syn_records)} {name})")

    # Sanity log: save the chosen ID list for reproducibility.
    ids_path = os.path.join(args.outdir, f"{prefix}.ids.txt")
    with open(ids_path, "w") as f:
        for bid in sampled:
            f.write(bid + "\n")
    print(f"Wrote {ids_path}")


if __name__ == "__main__":
    main()

"""
python select_500.py \
    --real       /home/jovyan/AURAD/detection/json/train_real.json \
    --cheff      /home/jovyan/AURAD/detection/json/train_cheff.json \
    --roentgen   /home/jovyan/AURAD/detection/json/train_roentgen.json \
    --roentgenv2 /home/jovyan/AURAD/detection/json/train_roentgenv2.json \
    --ours       /home/jovyan/AURAD/detection/json/train_ours.json \
    --n 500 \
    --seed 42 \
    --outdir .

"""