import os

import io

import base64

import numpy as np

import pandas as pd

import cv2

import matplotlib.pyplot as plt

import torch

import torch.nn as nn

from albumentations import Normalize

import time

from IPython.display import clear_output

from IPython.display import HTML

from loss_metric import dice_coef_metric_per_classes, jaccard_coef_metric_per_classes

def get_one_slice_data(img_name: str,

                       mask_name: str,

                       root_imgs_path: str = "images/",

                       root_masks_path: str = "masks/",) -> np.ndarray:

    img_path = os.path.join('images/', img_name)

    mask_path = os.path.join('masks/', mask_name)

    one_slice_img = cv2.imread(img_path)#[:,:,0] uncomment for grayscale

    one_slice_mask = cv2.imread(mask_path)

    one_slice_mask[one_slice_mask < 240] = 0  # remove artifacts

    one_slice_mask[one_slice_mask >= 240] = 255

    return one_slice_img, one_slice_mask

def get_id_predictions(net: nn.Module,

                       ct_scan_id_df: pd.DataFrame,

                       root_imgs_dir: str,

                       treshold: float = 0.3) -> list:

    """

    Factory for getting predictions and storing them and images in lists as uint8 images.

    Params:

        net: model for prediction.

        ct_scan_id_df: df with unique patient id.

        root_imgs_dir: root path for images.

        treshold: threshold for probabilities.

    """

    sigmoid = lambda x: 1 / (1 + np.exp(-x))

    images = []

    predictions = []

    net.eval()

    device = "cuda" if torch.cuda.is_available() else "cpu"

    print("device:", device)

    with torch.no_grad():

        for idx in range(len(ct_scan_id_df)):

            img_name = ct_scan_id_df.loc[idx, "ImageId"]

            path = os.path.join(root_imgs_dir, img_name)

            img_ = cv2.imread(path)

            img = Normalize().apply(img_)

            tensor = torch.FloatTensor(img).permute(2, 0, 1).unsqueeze(0)

            prediction = net.forward(tensor.to(device))

            prediction = prediction.cpu().detach().numpy()

            prediction = prediction.squeeze(0).transpose(1, 2, 0)

            prediction = sigmoid(prediction)

            prediction = (prediction >= treshold).astype(np.float32)

            predictions.append((prediction * 255).astype("uint8"))

            images.append(img_)

    return images, predictions

# Save image in original resolution

# helpful link - https://stackoverflow.com/questions/34768717/matplotlib-unable-to-save-image-in-same-resolution-as-original-image

def get_overlaid_masks_on_image(

                one_slice_image: np.ndarray,

                one_slice_mask: np.ndarray, 

                w: float = 512,

                h: float = 512, 

                dpi: float = 100,

                write: bool = False,

                path_to_save: str = '/content/',

                name_to_save: str = 'img_name'):

    """overlap masks on image and save this as a new image."""

    path_to_save_ = os.path.join(path_to_save, name_to_save)

    lung, heart, trachea = [one_slice_mask[:, :, i] for i in range(3)]

    figsize = (w / dpi), (h / dpi)

    fig = plt.figure(figsize=(figsize))

    fig.add_axes([0, 0, 1, 1])

    # image

    plt.imshow(one_slice_image, cmap="bone")

    # overlaying segmentation masks

    plt.imshow(np.ma.masked_where(lung == False, lung),

            cmap='cool', alpha=0.3)

    plt.imshow(np.ma.masked_where(heart == False, heart),

            cmap='autumn', alpha=0.3)

    plt.imshow(np.ma.masked_where(trachea == False, trachea),

               cmap='autumn_r', alpha=0.3) 

    plt.axis('off')

    fig.savefig(f"{path_to_save_}.png",bbox_inches='tight', 

                pad_inches=0.0, dpi=dpi,  format="png")

    if write:

        plt.close()

    else:

        plt.show()

def get_overlaid_masks_on_full_ctscan(ct_scan_id_df: pd.DataFrame, path_to_save: str):

    """

    Creating images with overlaid masks on each slice of CT scan.

    Params:

         ct_scan_id_df: df with unique patient id.

         path_to_save: path to save images.

    """

    num_slice = len(ct_scan_id_df)

    for slice_ in range(num_slice):

        img_name = ct_scan_id_df.loc[slice_, "ImageId"]

        mask_name = ct_scan_id_df.loc[slice_, "MaskId"]

        one_slice_img, one_slice_mask = get_one_slice_data(img_name, mask_name)

        get_overlaid_masks_on_image(one_slice_img,

                                one_slice_mask,

                                write=True, 

                                path_to_save=path_to_save,

                                name_to_save=str(slice_)

                                )

def create_video(path_to_imgs: str, video_name: str, framerate: int):

    """

    Create video from images.

    Params:

        path_to_imgs: path to dir with images.

        video_name: name for saving video.

        framerate: num frames per sec in video.

    """

    img_names = sorted(os.listdir(path_to_imgs), key=lambda x: int(x[:-4]))  # img_name must be numbers

    img_path = os.path.join(path_to_imgs, img_names[0])

    frame_width, frame_height, _ = cv2.imread(img_path).shape

    fourc = cv2.VideoWriter_fourcc(*'MP4V')

    video = cv2.VideoWriter(video_name + ".mp4", 

                            fourc, 

                            framerate, 

                            (frame_width, frame_height))

    for img_name in img_names:

        img_path = os.path.join(path_to_imgs, img_name)

        image = cv2.imread(img_path)

        video.write(image)

    cv2.destroyAllWindows()

    video.release()

def compute_scores_per_classes(model,

                               dataloader,

                               classes):

    """

    Compute Dice and Jaccard coefficients for each class.

    Params:

        model: neural net for make predictions.

        dataloader: dataset object to load data from.

        classes: list with classes.

        Returns: dictionaries with dice and jaccard coefficients for each class for each slice.

    """

    device = 'cuda' if torch.cuda.is_available() else 'cpu'

    dice_scores_per_classes = {key: list() for key in classes}

    iou_scores_per_classes = {key: list() for key in classes}

    with torch.no_grad():

        for i, (imgs, targets) in enumerate(dataloader):

            imgs, targets = imgs.to(device), targets.to(device)

            logits = model(imgs)

            logits = logits.detach().cpu().numpy()

            targets = targets.detach().cpu().numpy()

            dice_scores = dice_coef_metric_per_classes(logits, targets)

            iou_scores = jaccard_coef_metric_per_classes(logits, targets)

            for key in dice_scores.keys():

                dice_scores_per_classes[key].extend(dice_scores[key])

            for key in iou_scores.keys():

                iou_scores_per_classes[key].extend(iou_scores[key])

    return dice_scores_per_classes, iou_scores_per_classes