"""Patient-level split of the 484 MSD brain cases for the published autoencoder.

Cases are grouped into patients (exact re-uploads and other visits of one person) from
duplicate_pairs.csv. 40 patients with no case in 1-100 (the v1 autoencoder's training range, so both
models can be scored on the same unseen patients) are held out whole, one scan each for scoring.
Training keeps every other patient, minus exact re-uploads.

python split_patients.py duplicate_pairs.csv dataframe.csv <out dir>
"""

import csv
import random
import sys
from pathlib import Path

PATIENT, COPY, HELD_OUT = 0.6, 0.99, 40
pairs = [(r["case_a"], r["case_b"], max(float(r["flair_corr"]), float(r["t1w_corr"]))) for r in csv.DictReader(open(sys.argv[1]))]
frame = list(csv.DictReader(open(sys.argv[2])))
out = Path(sys.argv[3])
cases = sorted(r["subject"] for r in frame)


def components(threshold: float) -> dict[str, str]:
    parent = {c: c for c in cases}

    def root(c: str) -> str:
        while parent[c] != c:
            parent[c] = parent[parent[c]]
            c = parent[c]
        return c

    for a, b, corr in pairs:
        if corr > threshold:
            parent[root(a)] = root(b)
    return {c: root(c) for c in cases}


patient, scan = components(PATIENT), components(COPY)
patients: dict[str, list[str]] = {}
for c in cases:
    patients.setdefault(patient[c], []).append(c)
eligible = sorted((members for members in patients.values() if all(int(c[-3:]) > 100 for c in members)), key=lambda m: m[0])
held = random.Random(0).sample(eligible, HELD_OUT)
held_cases = {c for members in held for c in members}
# One copy per scan: the lowest-numbered case of each exact-duplicate set.
first_copy = {c for c in cases if c == min(o for o in cases if scan[o] == scan[c])}
train = [r for r in frame if r["subject"] not in held_cases and r["subject"] in first_copy]

with (out / "dataframe_train.csv").open("w", newline="") as f:
    writer = csv.DictWriter(f, fieldnames=list(frame[0]))
    writer.writeheader()
    writer.writerows(train)
(out / "heldout_cases.txt").write_text("\n".join(sorted(min(m) for m in held)) + "\n")
print(f"{len(patients)} patients; {len(cases) - len(first_copy)} exact re-uploads dropped")
print(f"held out {len(held)} patients ({len(held_cases)} scans, scored on one each); training {len(train)} scans")