"""PathGen-1.6M H&E-caption pairs (real text) -> cache/pathgen_meta.csv with a slide-level 80/10/10 split.
Used to bind the text modality to the H&E anchor (ImageBind: text only needs to be paired with images)."""
import json
import os

import numpy as np
import pandas as pd

from common import CACHE, ROOT

PG = os.path.join(ROOT, "..", "PathGen-1.6M")
d = json.load(open(os.path.join(PG, "Data", "PathGen-1.6M.json")))
m = pd.DataFrame([{"wsi_id": x["wsi_id"], "caption": x["caption"],
                   "image_path": os.path.join(PG, "Data", "images", f"{x['wsi_id']}_{x['position'][0]}_{x['position'][1]}.png")}
                  for x in d])
assert m.image_path.map(os.path.exists).all()
m["patient"] = m.wsi_id.str[:12]                         # TCGA-XX-XXXX: keep a patient's slides together
pats = np.random.default_rng(0).permutation(sorted(m.patient.unique()))
n = len(pats)
role = {p: "train" if i < 0.8 * n else "val" if i < 0.9 * n else "test" for i, p in enumerate(pats)}
m["split"] = m.patient.map(role)
m["project_code"] = m.wsi_id.str[5:7]
m.to_csv(os.path.join(CACHE, "pathgen_meta.csv"), index_label="idx")
print(m.split.value_counts().to_dict(), "patients:", n)