import os
import numpy as np
import random as rnd
import pandas as pd

import config
from utils import read_csv_column, read_csv, store_to_csv


def load_labels():
    return read_csv(config.PATIENT_LABELS_CSV)


def load_test_ids():
    return read_csv_column(config.TEST_PATIENTS_IDS)


def get_patient_name(patient_file):
    return os.path.basename(patient_file).split('.')[0]


def load_patient_ids():
    test_ids = set(load_test_ids())
    patient_ids = [get_patient_name(patient_id)
                   for patient_id in os.listdir(config.ALL_IMGS)
                   if get_patient_name(patient_id) not in test_ids]
                   
    return patient_ids

def get_class(labels, patient_id):
    return labels.get_value(patient_id, config.COLUMN_NAME)


def count_patients_from_class(patient_ids, labels, clazz):
    return len([patient for patient in patient_ids 
                if get_class(labels, patient) == clazz])


def split_data():
    labels = load_labels()
    total = len(labels)
    print("Total labels loaded: ", total)
    patient_ids = load_patient_ids()
    print("Total patient ids loaded: ", len(patient_ids))
    print("Patient with cancer are: ", count_patients_from_class(
        patient_ids, labels, config.CANCER_CLS))

    validation_size = int(0.15 * total)

    validation_set = rnd.sample(patient_ids, validation_size)
    train_set = [patient for patient in patient_ids 
                 if patient not in validation_set]

    print("Patients for training: ", len(train_set))
    print("Patients for validation: ", len(validation_set))

    print("Patients with cancer in validation set {}, no cancer {}.".format(
        count_patients_from_class(validation_set, labels, config.CANCER_CLS),
        count_patients_from_class(validation_set, labels, config.NO_CANCER_CLS)))
    print("Patients with cancer in training set {}, no cancer {}.".format(
        count_patients_from_class(train_set, labels, config.CANCER_CLS),
        count_patients_from_class(train_set, labels, config.NO_CANCER_CLS)))

    validation_labels = [get_class(labels, p) for p in validation_set]
    store_to_csv(validation_set, validation_labels, 
        config.VALIDATION_PATINETS_IDS)
    
    train_labels = [get_class(labels, p) for p in train_set]
    store_to_csv(train_set, train_labels, 
        config.TRAINING_PATIENTS_IDS)


if __name__ == '__main__':
    split_data()