"""Define evaluation splits. Each split file: splits/<name>/fold<k>.csv with columns tile_id, split.

  orion_official : the user's train/val/test_dataframe.csv (slide-disjoint)
  random_tile    : tile-level random 80/10/10 (slides shared across splits -> optimistic upper bound)
  patient_5fold  : 5-fold CV grouped by patient (CRC33_01/_02 kept together); val = 3 patients from train folds
  batch_lobo     : leave-one-IF-acquisition-batch-out (C59kX/mX/nX/qX); val = 2 patients from training batches
"""
import os

import numpy as np

from common import SPLITS, load_all_tiles

df = load_all_tiles()


def save(name, k, split):
    os.makedirs(os.path.join(SPLITS, name), exist_ok=True)
    out = df[["tile_id", "in_slide_name", "patient", "if_batch"]].copy()
    out["split"] = split
    out.to_csv(os.path.join(SPLITS, name, f"fold{k}.csv"), index=False)
    print(name, k, out.split.value_counts().to_dict(),
          {s: out[out.split == s].patient.nunique() for s in ["train", "val", "test"]})


# 1. official
save("orion_official", 0, df.orion_split.values)

# 2. random tile-level
rng = np.random.default_rng(0)
u = rng.random(len(df))
save("random_tile", 0, np.where(u < 0.8, "train", np.where(u < 0.9, "val", "test")))

# 3. patient-grouped 5-fold
pats = np.array(sorted(df.patient.unique()))
rng = np.random.default_rng(0)
perm = rng.permutation(pats)
folds = np.array_split(perm, 5)
for k in range(5):
    test_p = set(folds[k])
    rest = [p for p in perm if p not in test_p]
    val_p = set(rest[:3])
    split = np.where(df.patient.isin(test_p), "test", np.where(df.patient.isin(val_p), "val", "train"))
    save("patient_5fold", k, split)

# 4. leave-one-IF-batch-out
for k, b in enumerate(sorted(df.if_batch.unique())):
    rest = sorted(df[df.if_batch != b].patient.unique())
    val_p = set(np.random.default_rng(k).choice(rest, 2, replace=False))
    split = np.where(df.if_batch == b, "test", np.where(df.patient.isin(val_p), "val", "train"))
    save("batch_lobo", k, split)