import pandas as pd

import numpy as np

def get_splits(age_at_exam, patient_ids, exam_ids, splits, min_age_valid=16, max_age_valid=85, seed=0):

    rng = np.random.RandomState(seed)

    if sum(splits) > 1.0:

        raise ValueError('splits should be sum to a number smaller than one.')

    n_exams = len(exam_ids)

    # Get patients

    patients = np.unique(patient_ids)

    n_patients = len(patients)

    # Create correspondence

    hash_exams = dict(zip(exam_ids, range(n_exams)))

    hash_patients = dict(zip(patients, range(n_patients)))

    inverse_hash_patients = dict(zip(range(n_patients), patients))

    # Get all exams for each patient

    patient_exams = [[] for _ in range(n_patients)]

    for exam_idx in range(n_exams):

        exam_id = exam_ids[exam_idx]

        patient_id = patient_ids[exam_idx]

        patient_idx = hash_patients[patient_id]

        patient_exams[patient_idx].append(exam_id)

    # Get the age at one of the exams for each patient

    ages, _, _ = np.unique(age_at_exam, return_inverse=True, return_counts=True)

    patient_idx_per_age = {a: [] for a in ages}

    patient_single_exam = np.zeros(n_patients, dtype=int)

    for patient_idx in range(n_patients):

        # Pick random exam id for the given patient

        # OBS:Another formulation that could make sense could be to always pick the first exam....

        id_exam = rng.choice(patient_exams[patient_idx])

        patient_single_exam[patient_idx] = id_exam

        exam_idx = hash_exams[id_exam]

        a = age_at_exam[exam_idx]

        patient_idx_per_age[a].append(patient_idx)

    # Get number of patient in each split

    n_splits = [int(np.floor(s * n_patients)) for s in splits]

    n_splits += [n_patients - sum(n_splits)]

    # Shuffle

    rng.shuffle(ages)  # Shuffle ages

    for a, patient_idx in patient_idx_per_age.items():  # Shuffle within the same age

        rng.shuffle(patient_idx)

    # Pick one id per age and build a list from that

    all_patient_idx = []

    stop = False

    # Pick ids within the given range first (which will probably be used for training, validation and test)

    while not stop:

        stop = True

        for a in ages[(ages >= min_age_valid) & (ages <= max_age_valid)]:

            if patient_idx_per_age[a]:

                patient_idx = patient_idx_per_age[a].pop()

                all_patient_idx.append(patient_idx)

                stop = False

    # Pick remaining ids last (which will probably be used only for training)

    stop = False

    while not stop:

        stop = True

        for a in ages[(ages < min_age_valid) | (ages > max_age_valid)]:

            if patient_idx_per_age[a]:

                patient_idx = patient_idx_per_age[a].pop()

                all_patient_idx.append(patient_idx)

                stop = False

    # Save ids

    patients_in_splits = [[] for n in n_splits]

    single_exam_in_split = [[] for n in n_splits]

    exams_in_splits = [[] for n in n_splits]

    for i, patient_idx in enumerate(all_patient_idx):

        last_n = 0

        for s, n in enumerate(np.cumsum(n_splits)):

            if last_n <= i < n:

                patients_in_splits[s].append(inverse_hash_patients[patient_idx])

                single_exam_in_split[s].append(patient_single_exam[patient_idx])

                exams_in_splits[s] += patient_exams[patient_idx]

                last_n = n

    return patients_in_splits, single_exam_in_split, exams_in_splits

if __name__ == "__main__":

    import argparse

    import warnings

    parser = argparse.ArgumentParser(description='Generate data summary for the age prediction problem')

    parser.add_argument('file',

                        help='csv file to read data from.')

    parser.add_argument('--exam_id_col', default='N_exame',

                        help='column in csv containing exam id')

    parser.add_argument('--age_col', default='Idade',

                        help='column in csv containing age')

    parser.add_argument('--patient_id_col', default='N_paciente_univoco',

                        help='column in csv containing patient id')

    parser.add_argument('--splits', default=[0.15, 0.05], nargs='*', type=float,

                        help='percentage of data in each split')

    parser.add_argument('--splits_names', default=['test', 'val', 'train'], nargs='*', type=str,

                        help='split names')

    parser.add_argument('--no_plot', action='store_true',

                        help='dont show plots')

    args, unk = parser.parse_known_args()

    if unk:

        warnings.warn("Unknown arguments:" + str(unk) + ".")

    # Open csv file

    df = pd.read_csv(args.file, low_memory=False)

    # Remove duplicated rows

    df.drop_duplicates(args.exam_id_col, inplace=True)

    # Get ids from csv file

    exam_ids = np.array(df[args.exam_id_col], dtype=int)

    age_at_exam = np.array(df[args.age_col])

    patient_ids = np.array(df[args.patient_id_col], dtype=int)

    # define splits

    patients_in_splits, single_exam_in_split, exams_in_splits = get_splits(age_at_exam, patient_ids, exam_ids, args.splits)

    if not args.no_plot:

        import seaborn as sns

        import matplotlib.pyplot as plt

        n = len(args.splits) + 1

        fig, ax = plt.subplots(nrows=n)

        for i in range(n):

            age_single_exam = age_at_exam[np.isin(exam_ids, single_exam_in_split[i])]

            age = age_at_exam[np.isin(exam_ids, exams_in_splits[i])]

            sns.histplot(age, ax=ax[i], kde=False, bins=range(0, 130, 1))

            sns.histplot(age_single_exam, ax=ax[i], kde=False, bins=range(0, 130, 1))

        plt.show()