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

data_directory = "/home/ubuntu/thesis/data/ICIA2018/ICIAR2018_BACH_Challenge/WSI/"
#mask_list = os.listdir(data_directory + "masks/")
mask_list = ["A"+str(i+1).zfill(2)+".npy" for i in range(10)]
for m in tqdm(mask_list):
    scene = np.load(data_directory + m).transpose()
    image_slide = slideio.open_slide(data_directory + m[:-4] + ".svs")
    image_scene = image_slide.get_scene(0)

    dim0 = int(np.ceil(scene.shape[0] / 1024))
    dim1 = int(np.ceil(scene.shape[1] / 1024))
    resolutions = np.ceil(np.log2(max(dim0,dim1)))
    for r in range(int(resolutions) + 1):
        res = 2**r * 1024
        dim0 = int(np.ceil(scene.shape[0] / res))
        dim1 = int(np.ceil(scene.shape[1] / res))
        last_dim = (int(scene.shape[0] % res), int(scene.shape[1] % res))

        if last_dim[0] == 0 and last_dim[1] == 0:
            last_dim = (res, res)
        elif last_dim[0] == 0:
            last_dim = (res, last_dim[1])
        elif last_dim[1] == 0:  
            last_dim = (last_dim[0], res)

        for i in range(dim0):
            for j in range(dim1):
                if i == dim0-1 and j == dim1-1:
                    width = last_dim[0]
                    height = last_dim[1]
                elif i == dim0-1:
                    width = last_dim[0]
                    height = res
                elif j == dim1-1:
                    width = res
                    height = last_dim[1]
                else:
                    width = res
                    height = res
                #mask = scene.read_block((i*res,j*res, width, height), (width // (2**r), height // (2**r)))
                mask = scene[i*res:i*res+width:2**r, j*res:j*res+height:2**r]
                #mask = np.where(mask == 2, 1, 0).astype(np.uint8)
                if (np.max(mask) > 0):
                    image = image_scene.read_block((i*res,j*res, width, height), (width // (2**r), height // (2**r)))
                    # Save image and mask
                    # Save image
                    Image.fromarray(image).save(data_directory + "images_patches/" + m[:-9] + "_{}_{}_{}_{}_{}_{}.png".format(i*res,j*res,width, height, width // (2**r),  height // (2**r)))
                    # Save mask
                    Image.fromarray(mask).save(data_directory + "masks_patches/" + m[:-9] + "_{}_{}_{}_{}_{}_{}.png".format(i*res,j*res,width, height, width // (2**r),  height // (2**r)))