import numpy as np
from xml.dom import minidom
from skimage.draw import polygon
from tifffile import imread
import os
from PIL import Image

data_directory = "/home/ubuntu/thesis/data/MoNuSeg/MoNuSeg 2018 Training Data/"
def he_to_binary_mask(filename):
    im_file = data_directory + "Tissue Images/" + filename + '.tif'
    xml_file = data_directory + "Annotations/" + filename + '.xml'

    # Parse the XML file
    xDoc = minidom.parse(xml_file)
    Regions = xDoc.getElementsByTagName('Region')

    xy = []
    for regioni in range(Regions.length):
        Region = Regions.item(regioni)
        verticies = Region.getElementsByTagName('Vertex')
        xy_region = np.zeros((verticies.length, 2))
        for vertexi in range(verticies.length):
            x = float(verticies.item(vertexi).getAttribute('X'))
            y = float(verticies.item(vertexi).getAttribute('Y'))
            xy_region[vertexi] = [x, y]
        xy.append(xy_region)

    arr = imread(im_file)
    # Get image information
    im_info = {
        'Height': arr.shape[0],
        'Width': arr.shape[1]
    }
    binary_mask = np.zeros((im_info['Height'], im_info['Width']))
    color_mask = np.zeros((im_info['Height'], im_info['Width'], 3))

    for zz, region in enumerate(xy):
        print(f'Processing object # {zz + 1}')
        smaller_x = region[:, 0]
        smaller_y = region[:, 1]

        # Create binary and color masks
        polygon_mask = polygon(smaller_y, smaller_x, (im_info['Height'], im_info['Width']))
        binary_mask[polygon_mask] += zz + 1
        color_mask[polygon_mask] += np.random.rand(3)

    return binary_mask, color_mask

image_list = os.listdir(data_directory + "Tissue Images/")
for i in image_list:
    binary_mask, color_mask = he_to_binary_mask(i[:-4])
    values = np.unique(binary_mask)
    masks = np.zeros(binary_mask.shape)
    for k in range(len(values)):
        masks = np.where(binary_mask == values[k], k, masks)
    os.makedirs(data_directory + "Masks_new/", exist_ok=True)
    Image.fromarray(masks.astype(np.int32)).save(data_directory + "Masks_new/" + i[:-4] + ".png")