import os
import random
from os import listdir
from os.path import isfile, join, isdir


def get_image_paths(folder):
    image_paths = [join(folder, f) for f in listdir(folder) if isfile(join(folder, f)) and '.DS_Store' not in f]
    return image_paths


def get_alpha_name(number):
    alphas = "abcdefghijklmnopqrstuvwxyz"
    first = int(number / 26 / 26)
    second = int(number / 26) % 26
    third = number % 26
    return alphas[first] + alphas[second] + alphas[third]


# create an output folder if it does not already exist
def confirm_output_folder(output_folder):
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)


# parameters
input_folder = 'wsi_all_3x'
output_folder = 'wsi_3x'
train_split = 0.5
val_split = 0.15
test_split = 0.35

# get all the subfolders
all_input_folders = [join(input_folder, f) for f in listdir(input_folder) if isdir(join(input_folder, f))]

# get list of all images
all_image_paths = []
for folder in all_input_folders:
    all_image_paths += get_image_paths(folder)

# important dictionaries
image_path_to_pid = {}
pid_to_image_paths = {}

# fill the dictionaries
for image_path in all_image_paths:

    image_name = image_path.split('/')[-1]

    # parse the image_name for the patient id
    start = image_name.index('S')
    end = image_name.index('.')
    while image_name[end].lower() not in 'qwertyuiopasdfghjklzxcvbnm':
        end -= 1
    assert end > start
    pid = image_name[start:end]

    # add to dictionaries
    image_path_to_pid[image_path] = pid
    if pid in pid_to_image_paths:
        pid_to_image_paths[pid].append(image_path)
    else:
        pid_to_image_paths[pid] = [image_path]

print(len(image_path_to_pid), 'images from', len(pid_to_image_paths), 'patients')
patient_list = list(pid_to_image_paths.keys())
random.seed(0)
random.shuffle(patient_list)
num_patients = len(patient_list)

# partition data
train_patients = patient_list[:int(train_split*num_patients)]
val_patients = patient_list[int(train_split*num_patients):int((train_split+val_split)*num_patients)]
test_patients = patient_list[int((train_split+val_split)*num_patients):]

train_images = []
for pid in train_patients:
    train_images += pid_to_image_paths[pid]

val_images = []
for pid in val_patients:
    val_images += pid_to_image_paths[pid]

test_images = []
for pid in test_patients:
    test_images += pid_to_image_paths[pid]

print('train_patients:', len(train_patients), 'with', len(train_images), 'images')
print('val_patients:', len(val_patients), 'with', len(val_images), 'images')
print('test_patients:', len(test_patients), 'with', len(test_images), 'images')

# encode image names
all_image_paths_ordered = train_images + val_images + test_images
image_path_to_code = {}
for i, image_path in enumerate(all_image_paths_ordered):
    image_path_to_code[image_path] = get_alpha_name(i)

folder_to_count = {}

writer = open(join(input_folder, 'image_to_pid.tsv'), 'w')
# copy all the images into the new locations!
for image_set, set_type in zip([train_images, val_images, test_images], ['train', 'val', 'test']):
    for image_path in image_set:
        label = image_path.split('/')[-2]

        folder = '/'.join([output_folder, set_type, label])

        if folder in folder_to_count:
            folder_to_count[folder] += 1
        else:
            folder_to_count[folder] = 1

        new_path = '/'.join([output_folder, set_type, label, image_path_to_code[image_path]+'.jpg'])

        write_line = new_path + '\t' + image_path
        writer.write(write_line + '\n')

print(folder_to_count)

test_img_file = open('wsi_all_3x/new_test_images.txt', 'r').readlines()
test_imgs = [x[:-1] for x in test_img_file]
test_patients = []

img_to_pid = {image_path_to_code[img_path]: image_path_to_pid[img_path] for img_path in image_path_to_pid}

for test_img in test_imgs:
    test_patients.append(img_to_pid[test_img[:3]])

print(len(set(test_patients)))