from core.datasets.BaseDataset import BaseDataset
from core.explanations.utils import (
    GREEN, GRAY, RED,
    interp_color
)
from matplotlib import pyplot as plt
from skimage.color import label2rgb

import numpy as np


class ImageDataset(BaseDataset):

    def __init__(self, name, num_classes, num_features, sample_shape, seed):
        super().__init__(
            name=name,
            path=f'models/image/{name}.onnx',
            num_classes=num_classes,
            num_features=num_features,
            sample_shape=sample_shape,
            seed=seed
        )

    def plot_explanation(self, sample, explanation, figure_name):
        plot_shape = self.sample_shape[:2]

        # Color the pixels in the explanation
        mask = np.zeros(self.num_features, dtype=bool)
        mask[list(explanation)] = True

        overlay = label2rgb(
            mask.reshape(plot_shape),
            image=sample.reshape(plot_shape),
            colors=[[0, 1, 0]],
            bg_label=0,
            saturation=1
        )

        figure, ax = plt.subplots()
        ax.imshow(overlay)

        self.__plot(figure, ax, figure_name)

    def plot_feature_importance(self, feature_importance, figure_name):
        figure, ax = plt.subplots()
        ax.set_xlabel("Features")
        ax.set_ylabel("Feature Importance")

        # Color bars based on importance
        importance = [feature_importance[i] for i in range(self.num_features)]
        importance = (importance - np.min(importance)) / (np.max(importance) - np.min(importance))

        # Interpolate from GRAY to RED
        colors = [interp_color(GRAY, RED, imp) for imp in importance]

        self.__plot(figure, ax, figure_name)

    def __plot(self, figure, ax, figure_name):
        ax.axis('off')
        self._save_figure(figure, figure_name)
