from core.datasets.BaseDataset import BaseDataset
from core.explanations.utils import (
    GREEN, GRAY, RED,
    interp_color
)

import matplotlib.patches as mpatches
import matplotlib.pyplot as plt
import numpy as np


class CustomDataset(BaseDataset):

    def __init__(self, name, num_classes, num_features, sample_shape, seed):
        super().__init__(
            name=name,
            path=f'models/custom/{name}.onnx',
            num_classes=num_classes,
            num_features=num_features,
            sample_shape=sample_shape,
            seed=seed
        )
        self.feature_names = [n for n in range(num_features)]

    def plot_explanation(self, sample, explanation, figure_name):
        figure, ax = plt.subplots()
        ax.set_xlabel("Features")
        ax.set_ylabel("Feature Value")

        # Color bars that are in the explanation
        colors = []
        for tick_label, feature in zip(ax.get_xticklabels(), range(self.num_features)):
            if feature in explanation:
                color = GREEN
                colors.append(color)
                tick_label.set_color(color)
            else:
                color = GRAY
                colors.append(color)
                tick_label.set_color(color)

        # Add a legend to indicate relevant and irrelevant features
        legend_patches = [
            mpatches.Patch(color=GREEN, label='Relevant Feature'),
            mpatches.Patch(color=GRAY, label='Irrelevant Feature')
        ]
        ax.legend(handles=legend_patches)

        self.__plot(figure, ax, sample, colors, 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, importance, colors, figure_name)

    def __plot(self, figure, ax, y, colors, figure_name):
        ax.bar(self.feature_names, y, color=colors)

        # Set the x-ticks to the feature names and rotate them for better visibility
        x_positions = np.arange(self.num_features)
        ax.bar(x_positions, y, color=colors)
        ax.set_xticks(x_positions)

        self._save_figure(figure, figure_name)
