import cv2
from PIL import Image
import numpy as np, torchvision
import torch, os, json
from skimage.color import label2rgb
from PIL import ImageColor
from cxas.label_mapper import label_mapper, id2label_dict, colors, colors_alpha, category_colors, category_ids
from cxas.file_io import FileLoader



def get_img(
            path: str
           ) -> np.array:
    """
    Load image from path.

    Parameters
    ----------
        path: Image file path
        
    Returns
    ----------
        image: Image in numpy format

    """
    assert os.path.isfile(path)
    loader  = FileLoader('')
    return loader.load_file(path)['orig_data']

def get_label(
            path: str
            ) -> np.array:
    """
    Load label from path with .npy-suffix.

    Parameters
    ----------
        path: label file path
        
    Returns
    ----------
        mask: Label mask in numpy format

    """
    assert os.path.isfile(path)
    return np.load(path)

def visualize_label(label: np.array, 
                    img: np.array, 
                    label_to_visualize: np.array, 
                    concat:bool = False, 
                    axis:int = 1
                   ) -> Image:
    """
    visualize certain labels from mask

    Parameters
    ----------
        label:  Mask in shape [classes (159), width, height]
        img:    Image in shape [classes (159), width, height]
        concat: Whether to display image and visualization side by side
        axis:   axis at which image and visualization are shown side by side
        
    Returns
    ----------
        visualization: Label visualization as PIL Image
        
    """
    

    out_mask = np.zeros((img.shape[0],img.shape[1],3)).astype(np.uint8)

    for i in label_to_visualize:
        imgray = (label[i,:,:]*255).astype(np.uint8)
        ret, thresh = cv2.threshold(imgray, 127, 255, 0)
        x = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
        for contour in x:
            if (contour is not None) and (len(contour) > 0) and  (len(contour[0])>2) and (type(contour) == type(())):
                cv2.fillPoly(
                             out_mask, 
                             contour, 
                             # [int(j*255) for j in colors[i]],
                            [colors_alpha[i][0]*255,colors_alpha[i][1]*255,colors_alpha[i][2]*255,colors_alpha[i][3]],
                            )
        
    out_contour = np.zeros((img.shape[0],img.shape[1],3)).astype(np.uint8)
    for i in label_to_visualize:
        imgray = (label[i,:,:]*255).astype(np.uint8)
        ret, thresh = cv2.threshold(imgray, 127, 255, 0)
        x = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
        for contour in x:
            if (contour is not None) and (len(contour) > 0) and  (len(contour[0])>2)  and (type(contour) == type(())):
                cv2.drawContours(out_contour, contour, -1, [int(j*255) for j in colors[i]], 2)
    
    out = cv2.addWeighted(out_contour, 1, out_mask, 0.7, 0.0)
    out = cv2.addWeighted(img, 0.5, out, 0.7, 0.0)
    
    if concat:
        return Image.fromarray(np.concatenate([img, out],axis)).convert('RGB')
    else:
        return Image.fromarray(out).convert('RGB')

def visualize_mask(class_names: list, 
                   mask:np.array,
                   image:np.array,
                   img_size: int,
                   cat:bool=True,
                   axis:int=1,
                  ) -> Image:
    """
    Resize image and label to desired size and visualize certain labels

    Parameters
    ----------
        class_names: List of classes of interest
        mask:  Mask in shape [classes (159), width, height]
        image:    Image in shape [3, width, height]
        img_size: Desired image size
        cat: Whether to display image and visualization side by side
        axis:   axis at which image and visualization are shown side by side
        
    Returns
    ----------
        visualization: Label visualization as PIL Image in desired size
        
    """
    # Resize image and label to desired size
    img = torch.nn.functional.interpolate(
                torch.tensor(image).float().unsqueeze(0),
                img_size,
                mode='bilinear'
            ).byte().numpy()[0]
        
        # reshape to match cv2 shapes
    img = np.transpose(img, [1,2,0])
    
    label = torch.nn.functional.interpolate(
                torch.tensor(mask).float().unsqueeze(0),
                img_size,
                mode='nearest'
            ).bool().numpy()[0]

    
    if type(class_names) == list:
        pass
    else:
        class_names = [class_names]
    
    label_to_visualize = np.concatenate([np.array(label_mapper[n]).flatten() for n in class_names]).flatten()
    
    return visualize_label(label, img,  label_to_visualize, cat, axis)

def visualize_from_file(class_names: list, 
                         img_path: str, 
                         label_path: str, 
                         img_size: int, 
                         cat:bool = True, 
                         axis:int = 1, 
                         do_store:bool = False,
                         out_dir:str = '',
                        ) -> Image:
    """
    Load Image and label, resize image and label to desired size, and visualize certain labels

    Parameters
    ----------
        class_names: list of class names to visualize
        img_path: path of the image file to visualize
        label_path: path of the label file to visualize
        img_size: size at which the image should be visualized
        cat: show image and label side by side
        axis: axis at which image and label are shown side by side
        do_store: boolean indicating whether to store the visualization in the out_dir 
                            with the associated label path and class_name
        out_dir: path at which to store visualization
        
    Returns
    ----------
        visualization: Pillow Image with labels overlaying original image
        
    """

    # Load image and label files
    img = get_img(img_path)
    label = get_label(label_path)
    # visualize desired labels
    visualization = visualize_mask(class_names, label, img, img_size, cat, axis)
    if do_store:
        os.makedirs(os.path.dirname(out_dir), exist_ok=True)
        visualization.save(os.path.join(
            out_dir,
            '{}_{}.png'.format(os.path.basename(label_path).split('.')[0], '_'.join(class_names))))
    return visualization