#
#   Lightnet dataset that works with brambox annotations
#   Copyright EAVISE
#
# https://eavise.gitlab.io/lightnet/_modules/lightnet/models/_dataset_brambox.html#BramboxDataset
# https://eavise.gitlab.io/brambox/notes/02-getting_started.html#Loading-data

import os
import copy
import logging
from PIL import Image
import numpy as np
import lightnet.data as lnd
from pathflowai.utils import load_sql_df
import dask.array as da
from os.path import join

try:
    import brambox as bb
except ImportError:
    bb = None

__all__ = ['BramboxDataset']
log = logging.getLogger(__name__)

# ADD IMAGE ANNOTATION TRANSFORM
# ADD TRAIN VAL TEST INFO

class BramboxPathFlowDataset(lnd.Dataset):
    """ Dataset for any brambox annotations.

    Args:
        annotations (dataframe): Dataframe containing brambox annotations
        input_dimension (tuple): (width,height) tuple with default dimensions of the network
        class_label_map (list): List of class_labels
        identify (function, optional): Lambda/function to get image based of annotation filename or image id; Default **replace/add .png extension to filename/id**
        img_transform (torchvision.transforms.Compose): Transforms to perform on the images
        anno_transform (torchvision.transforms.Compose): Transforms to perform on the annotations

    Note:
        This dataset opens images with the Pillow library
    """
    def __init__(self, input_dir, patch_info_file, patch_size, annotations, input_dimension, class_label_map=None, identify=None, img_transform=None, anno_transform=None):
        if bb is None:
            raise ImportError('Brambox needs to be installed to use this dataset')
        super().__init__(input_dimension)

        self.annos = annotations
        self.annos['ignore']=0
        self.annos['class_label']=self.annos['class_label'].astype(int)#-1
        print(self.annos['class_label'].unique())
        #print(self.annos.shape)
        self.keys = self.annos.image.cat.categories # stores unique patches
        #print(self.keys)
        self.img_tf = img_transform
        self.anno_tf = anno_transform
        self.patch_info=load_sql_df(patch_info_file, patch_size)
        IDs=self.patch_info['ID'].unique()
        self.slides = {slide:da.from_zarr(join(input_dir,'{}.zarr'.format(slide))) for slide in IDs}
        self.id = lambda k: k.split('/')
        # experiment
        #self.annos['x_top_left'], self.annos['y_top_left']=self.annos['y_top_left'], self.annos['x_top_left']
        self.annos['width'], self.annos['height']=self.annos['height'], self.annos['width']
        # Add class_ids
        if class_label_map is None:
            log.warning(f'No class_label_map given, generating it by sorting unique class labels from data alphabetically, which is not always deterministic behaviour')
            class_label_map = list(np.sort(self.annos.class_label.unique()))
        self.annos['class_id'] = self.annos.class_label.map(dict((l, i) for i, l in enumerate(class_label_map)))

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

    @lnd.Dataset.resize_getitem
    def __getitem__(self, index):
        """ Get transformed image and annotations based of the index of ``self.keys``

        Args:
            index (int): index of the ``self.keys`` list containing all the image identifiers of the dataset.

        Returns:
            tuple: (transformed image, list of transformed brambox boxes)
        """
        if index >= len(self):
            raise IndexError(f'list index out of range [{index}/{len(self)-1}]')

        # Load
        #print(self.keys[index])
        ID,x,y,patch_size=self.id(self.keys[index])
        x,y,patch_size=int(x),int(y),int(patch_size)
        img = self.slides[ID][x:x+patch_size,y:y+patch_size].compute()#Image.open(self.id(self.keys[index]))
        anno = bb.util.select_images(self.annos, [self.keys[index]])

        # Transform
        if self.img_tf is not None:
            img = self.img_tf(img)
        if self.anno_tf is not None:
            anno = self.anno_tf(anno)

        return img, anno