import os
import shutil
import numpy as np
import slideio
from PIL import Image
from tqdm import tqdm

data_directory = '/home/ubuntu/thesis/data/MoNuSAC/'
image_source = data_directory + 'MoNuSAC_images_and_annotations/'
mask_source = data_directory + 'MoNuSAC_masks/'
image_destination = data_directory + 'images/'
mask_destination = data_directory + 'masks/'
image_list = os.listdir(image_source)

# Training data
for i in image_list:
    subimage_list = os.listdir(image_source + i)
    for j in subimage_list:
        if j.endswith('.tif'):
            target_image_file = image_destination + j
            shutil.copyfile(image_source + i + '/' + j, target_image_file)
            types = ["Epithelial", "Lymphocyte", "Macrophage", "Neutrophil"]
            for t in types:
                if os.path.exists(mask_source + i + '/' + j[:-4] + '/' + t):
                    mask_list = os.listdir(mask_source + i + '/' + j[:-4] + '/' + t)
                    for m in mask_list:
                        target_mask_file = mask_destination + j[:-4] + '_' + t + '.png'
                        slide = slideio.open_slide(mask_source + i + '/' + j[:-4] + '/' + t + '/' + m)
                        scene = slide.get_scene(0)
                        mask = scene.read_block((0,0, scene.size[0], scene.size[1]))
                        Image.fromarray(mask.astype(np.uint8)).save(target_mask_file)
                        
# Testing data
source = data_directory + 'MoNuSAC_Testing_Color_Coded_Masks/'
image_list = os.listdir(source)
for i in tqdm(image_list):
    subimage_list = os.listdir(source + i)
    for j in subimage_list:
        if j.endswith('.png'):
            target_image_file = image_destination + j
            shutil.copyfile(source + i + '/' + j, target_image_file)
        else:
            target_mask_file = mask_destination + j[:-17] + '.png'
            mask_array = np.array(Image.open(source + i + '/' + j))
            binary_arr = np.zeros((mask_array.shape[0], mask_array.shape[1]))
            for x in range(mask_array.shape[0]):
                for y in range(mask_array.shape[1]):
                    if mask_array[x, y, 0] == 255 or mask_array[x, y, 1] == 255 or mask_array[x, y, 2] == 255:
                        binary_arr[x][y] = 1
            Image.fromarray(binary_arr.astype(np.uint8)).save(target_mask_file)