# %%
from collections import defaultdict
import json
import os

import numpy as np
import pandas as pd

from matplotlib.pyplot import hist
from skmultilearn.model_selection import iterative_train_test_split

# 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
    display_sample,
    get_labeled,
    get_y_Path,
    make_dcmdicts,
    path_2dcm_int16,
    path_2label,
)
from adpkd_segmentation.data.data_utils import (  # noqa
    PATIENT,
    SEQUENCE,
    KIDNEY_PIXELS,
    MR,
    VOXEL_VOLUME,
)

STUDY_TKV = "study_tkv"

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


# TKV checks
# %%
def TKV_update(dcm2attribs):
    studies = defaultdict(int)
    for dcm, attribs in dcm2attribs.items():
        study_id = (attribs[PATIENT], attribs[MR])
        studies[study_id] += attribs[KIDNEY_PIXELS] * attribs[VOXEL_VOLUME]

    for dcm, attribs in dcm2attribs.items():
        tkv = studies[(attribs[PATIENT], attribs[MR])]
        attribs[STUDY_TKV] = tkv

    return studies, dcm2attribs


# %%
studies, dcm2attribs = TKV_update(dcm2attribs)
hist(studies.values(), bins=40)

# %%
hist(np.log(list(studies.values())), bins=40)

# %%
# check outliers
outliers = []
for study, tkv in studies.items():
    if np.log(tkv) < 11:
        outliers.append((study, tkv))

print(outliers)

# %%
images = []
labels = []
sequences = set()
for dcm_path, attribs in dcm2attribs.items():
    study_id = (attribs[PATIENT], attribs[MR])
    if study_id == outliers[0][0] and attribs[KIDNEY_PIXELS] > 0:
        im = path_2dcm_int16(dcm_path)
        label = path_2label(get_y_Path(dcm_path))
        images.append(im)
        labels.append(label)
        sequences.add(attribs[SEQUENCE])

# %%
for im, label in zip(images, labels):
    display_sample((im, label))

# %%
print("Outlier sequence: ", sequences)

# %%
# Patient info

patient_info = set()
for dcm_path, attribs in dcm2attribs.items():
    patient = attribs[PATIENT]
    seq = attribs[SEQUENCE]
    tkv = attribs[STUDY_TKV]
    mr = attribs[MR]
    patient_info.add((patient, seq, mr, tkv))

print(patient_info)


# %%
df = pd.DataFrame.from_records(
    list(patient_info),
    columns=[PATIENT, SEQUENCE, MR, STUDY_TKV],
    index=PATIENT,
).sort_index()

# %%
df.to_csv("./notebooks/patients_2020_09_06.csv")

# %%
print(df.index.value_counts())

# %%
print(df.seq.value_counts())

# %%
print(df.study_tkv.describe())

print(np.log(df.study_tkv).describe())

# %%
# patient IDs, wtihout the outlier
outlier_ids = {out[0][0] for out in outliers}
print(outlier_ids)
all_ids = [id_ for id_ in all_patient_IDS if id_ not in outlier_ids]


# %%
def create_label_arrays(patient_info, all_ids):
    patient_to_label = {}
    for id_ in all_ids:
        # previously 7
        # SSFSE, FIESTA, OTHER, LOG_TKV_1, LOG_TKV_2,
        # LOG_TKV_3, LOG_TKV_4
        patient_to_label[id_] = np.zeros(6, dtype=np.uint8)
    for patient, seq, mr, tkv in patient_info:
        # outlier
        if patient not in patient_to_label:
            continue
        # sequence labeling
        # the same patient can have more
        if "SSFSE" in seq:
            patient_to_label[patient][0] = 1
        elif "FIESTA" in seq:
            patient_to_label[patient][1] = 1
        else:
            # previously a separate label
            patient_to_label[patient][1] = 1
        # LOG TKV category
        # old values:
        # 13.6 to 15.1 interquartile range
        # 14.2 median
        log_tkv = np.log(tkv)
        if log_tkv < 13.7:
            patient_to_label[patient][2] = 1
        elif 13.7 <= log_tkv < 14.2:
            patient_to_label[patient][3] = 1
        elif 14.2 <= log_tkv < 14.8:
            patient_to_label[patient][4] = 1
        elif 14.8 <= log_tkv:
            patient_to_label[patient][5] = 1

    return patient_to_label


# %%
patient2label = create_label_arrays(patient_info, all_ids)

# %%
labels = [patient2label[id_] for id_ in all_ids]
X = np.array(all_ids)[..., np.newaxis]
y = np.stack(labels)

print(y.shape)
print(len(all_ids))
print(X.shape)

# %%
# Split to train, val, test
TRAIN = 0.7
VAL = 0.15
TEST = 0.15
np.random.seed = 42

X_train_val, y_train_val, X_test, y_test = iterative_train_test_split(
    X, y, test_size=TEST
)

X_train, y_train, X_val, y_val = iterative_train_test_split(
    X_train_val, y_train_val, test_size=VAL / (TRAIN + VAL)
)

# %%
print("Sizes: ", len(X_train), len(X_val), len(X_test))

# %%
print(y_train)
print(y_val)
print(y_test)

# %%
print(df[df.index.isin(X_test.squeeze())])
print(df[df.index.isin(X_val.squeeze())])

# %%
split_dict = {
    "train": list(X_train.squeeze()),
    "val": list(X_val.squeeze()),
    "test": list(X_test.squeeze()),
}

# %%
with open("./stratification/strat_split_2020_09_06.json", "w") as f:
    json.dump(split_dict, f, indent=4)

# %%