from core.datasets.BaseDataset import BaseDataset
from core.explanations.utils import (
    GRAY, GREEN, RED, BLUE,
    interp_color,
    group_samples_by_label
)
from matplotlib import pyplot as plt

import matplotlib.patches as mpatches
import numpy as np


class ClassicDataset(BaseDataset):

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

    def get_samples(self, label, num_samples):
        (x_train, y_train), (x_test, y_test) = self.get_data()
        x = np.concatenate([x_train, x_test], axis=0)
        y = np.concatenate([y_train, y_test], axis=0)

        samples_by_label = group_samples_by_label(x, y, self.num_classes)
        return samples_by_label[label][:num_samples]

    def plot_explanation(self, sample, explanation, bounds, figure_name):
        figure, ax = plt.subplots(figsize=(12, 6))
        ax.set_xlabel("Features")
        ax.set_ylabel("Feature Value")
        ax.set_ylim(0, 1)

        # Create positions for features
        x_positions = np.arange(self.num_features)

        # Plot rectangles for each feature spanning from lower to upper bound
        for i in range(self.num_features):
            if i < len(bounds):
                lower_bound, upper_bound = bounds[i]
                color = GREEN if i in explanation else BLUE

                # Create rectangle from lower_bound to upper_bound at feature position
                rect_width = 0.6  # Width of the rectangle
                rect = plt.Rectangle((i - rect_width/2, lower_bound), rect_width, upper_bound - lower_bound,
                                   color=color, alpha=0.7)
                ax.add_patch(rect)
            else:
                # If no bound provided for this feature, skip
                continue

        # Set x-ticks and labels
        ax.set_xticks(x_positions)
        ax.set_xticklabels(self.feature_names, rotation=45, ha='right')

        # Add margins to prevent first and last features from being cut off
        ax.set_xlim(-0.5, self.num_features - 0.5)

        # Color tick labels based on whether feature is in explanation
        for i, tick_label in enumerate(ax.get_xticklabels()):
            tick_label.set_color('black')

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

        self._save_figure(figure, 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)
        ax.set_xticklabels(self.feature_names, rotation=45, ha='right')

        self._save_figure(figure, figure_name)
