# %%
import json
import os
import random

# enable lib loading even if not installed as a pip package or in PYTHONPATH
# also convenient for relative paths in example config files
from pathlib import Path

os.chdir(Path(__file__).resolve().parent.parent)

from adpkd_segmentation.data.data_utils import ( # noqa
    get_labeled,
    make_dcmdicts,
)
from adpkd_segmentation.data.link_data import makelinks  # noqa

# %%
makelinks()
# %%
dcm_paths = sorted(get_labeled())
dcm2attribs, patient2dcm = make_dcmdicts(tuple(dcm_paths))
all_patient_IDS = list(patient2dcm.keys())


# %%
with open("./stratification/strat_split_2020_09_24.json", "r") as f:
    full_split = json.load(f)

# %%
new_train = [
    patient_id
    for patient_id in all_patient_IDS
    if patient_id not in full_split["val"]
    and patient_id not in full_split["test"]
]

# %%
random.seed(42)
random.shuffle(new_train)

# %%
new = {}
new["val"] = full_split["val"]
new["test"] = full_split["test"]
new["train"] = new_train
# %%
with open("./stratification/strat_split_2020_09_28_extended.json", "w") as f:
    json.dump(new, f, indent=4)

# %%