import os
import numpy as np
from PIL import Image 




def generate_checkerboard(img_width, img_height, grid_size):
    ck = np.kron([[255, 0] * int(grid_size/2), [0, 255] * int(grid_size/2)]*int(grid_size/ 2) , np.ones((int(img_height / grid_size), int(img_width/ grid_size))))
    
    #ck = np.expand_dims(ck, axis=2)

    ck = ck.astype(np.uint8) # this is needed ot use ToTensor() transformaion in pytorch

    return ck


def get_tiled_ground_truth(mask, grid_size):
    height = mask.shape[0]
    width = mask.shape[1]
    tile_height = height / grid_size
    tile_width = width / grid_size
    mask = mask[:, : , 0] # get only 0th channel
    
    tile_mask = np.zeros_like(mask, dtype=np.uint8)

   # if grid_size == mask.shape[0]:
    #    return mask
    
    for c in range(grid_size):
        
        for r in range(grid_size):
        
            #mask = (mask > 128) # convert mask to binary
            #print(mask)

            row_start = int(r*tile_height)
            row_end = int(r*tile_height + tile_height)
            column_start =  int(c * tile_width) 
            column_end = int(c * tile_width + tile_width)
            
            #print("row start=", row_start)
            #print("row end=", row_end)
            #print("column start=", column_start)
            #print("colum end=", column_end)

            tile_sum = np.sum(mask[row_start:row_end, column_start:column_end])

            #print(tile_sum)  
            
            if tile_sum > 0:
                tile_mask[row_start:row_end, column_start:column_end] = 255
            
        #tile = 
    tile_mask = tile_mask.astype(np.uint8)
    #plt.imshow(tile_mask)
    return tile_mask




class DatasetWithGridEncoding(object):
    """
    A class of Dataset whch is supported for Pytorch.

    ...

    Attributes
    ----------
    imgs_root: str
        A root directory contains image files

    masks_root: str
        A root directory contains mask files. Please use the image names as prefix for mask's names to find correct mask for correct image.

    img_size: int
        A size of image. Image width and height are resized into this given size.

    grid_sizes: int array
        A array of int sizes which is used as grid sizes (default is  [2,4,8,16,32,64,128,256])

    transforms: Pytorch transforms
        Other type of transforms which come with Pytorch library. 

    Methods
    -------
    __getitem__(idx):
        Get an single data item from the dataset with image, gird_encoding image and correspondig mask.
    """

    def __init__(self, imgs_root, masks_root, img_size = 256, grid_sizes=[2,4,8,16,32,64,128,256] , transforms = None):
        """
        Parameters
        ----------
        imgs_root: str
            A root directory contains image files

        masks_root: str
            A root directory contains mask files. Please use the image names as prefix for mask's names to find correct mask for correct image.

        img_size: int
            A size of image. Image width and height are resized into this given size.

        grid_sizes: int array
            A array of int sizes which is used as grid sizes (default is  [2,4,8,16,32,64,128,256])

        transforms: Pytorch transforms
            Other type of transforms which come with Pytorch library.
        """

        self.imgs_root = imgs_root
        self.masks_root = masks_root

        self.img_size = img_size

        self.transforms = transforms

        self.length_of_grid_sizes = len(grid_sizes)

        self.imgs = list(sorted(os.listdir(self.imgs_root)))
        self.masks = list(sorted(os.listdir(self.masks_root)))

        self.num_imgs = len(self.imgs)

        self.imgs = self.imgs * self.length_of_grid_sizes
        self.masks = self.masks * self.length_of_grid_sizes

        self.grid_sizes_repeated = np.repeat(grid_sizes, self.num_imgs)

        self.all_in_one = list(zip(self.imgs, self.masks, self.grid_sizes_repeated)) #(img, mask, grid_size)

       # self.imgs_repeated
        #self.masks_repea

    def __len__(self):
        return len(self.all_in_one)


    def __getitem__(self, idx):
        """ Get an single data item from the dataset with image, gird_encoding image and correspondig mask. 
        
        Main function which return data for Pytorch dataloader.

        Parameters
        ---------
        idx: index of the datarecord to be returned. 

        Returns
        -------
        dictionary
            A dictionary with three keys and corresponding values, {"img":img, "grid_encode": grid_encode, "mask":mask}
        """

        img_path = os.path.join(self.imgs_root, self.all_in_one[idx][0]) # 0th one= image
        mask_path = os.path.join(self.masks_root, self.all_in_one[idx][1]) # 1st one = mask
        grid_size = self.all_in_one[idx][2] # 2nd one = grid size

        img_size = self.img_size

        img = Image.open(img_path)
        mask = Image.open(mask_path)

        grid_encode = generate_checkerboard(img_size, img_size, grid_size)#[:, :, 0]
    

        # resizing
        img = img.resize((img_size, img_size), Image.NEAREST)
        mask = mask.resize((img_size, img_size), Image.NEAREST)

        # covert to numpy

        img = np.array(img)
        mask = np.array(mask) #



        # clean mask values (this is done to remove small values in mask images)
        mask = (mask > 128) * 255 # if mask value > 128, set it to 255 (this will remove 254, 253 values and 1,2 etc)
         
        mask = get_tiled_ground_truth(mask, grid_size)
        #mask = mask[:, :, 0]

       

        if self.transforms is not None:
            # img, target= self.transforms(img, target)

            img = self.transforms(img)
            mask = self.transforms(mask)
            grid_encode = self.transforms(grid_encode)

          
        return {"img":img, "grid_encode": grid_encode, "mask":mask}