from core.models.BaseModel import BaseModel
from scipy.ndimage import convolve

import numpy as np
import random


class OrderingModel(BaseModel):
    sharpness = 10

    def __init__(self, dataset, label, seed, sharpness=10):
        super().__init__(dataset, seed)
        self.label = label
        self.sharpness = sharpness

    def traversal_order(self, sample, ordering):
        """
        Compute the order in which to check if the features are relevant for the model's prediction.
        There are three possible orders:
          1. 'linear': sorts the features in a linear order.
          2. 'random': shuffles the features randomly.
          3. 'verix': sorts the features based on the Verix method.
          4. 'sensitivity': sorts the features based on the model's output sensitivity to each feature and applies a smoothing kernel to account for positional weighting.

        :param sample: the sample to explain.
        :param ordering: 'linear', 'random', 'verix', 'sensitivity' (by default).
        :return: outputs the ordered features.
        """
        features = np.arange(np.prod(sample.shape))

        if ordering == "linear":
            features.sort()
        elif ordering == "random":
            random.shuffle(features)
        elif ordering == "verix":
            sensitivity = self.__get_sensitivity(sample)
            features = sensitivity.argsort()
        elif ordering == "sensitivity":
            if sample.ndim == 1:
                sensitivity = self.__get_sensitivity(sample)
            else:
                # Apply a smoothing kernel to the sensitivity to account for the positional weighting.
                sensitivity = self.__get_smooth_sensitivity(sample)

            features = sensitivity.argsort()
        else:
            print("Traversal not supported: try 'random', 'linear', 'verix', or 'sensitivity'.")

        return features

    def __get_sensitivity(self, sample):
        # Get the original prediction.
        pred_orig = self.session.run(None, {self.model.graph.input[0].name: sample.reshape(1, *sample.shape)})[0]

        # Get the sensitivity of each feature.
        sensitivities = []

        for i in range(self.dataset.num_features):
            # Perturb the sample by flipping the i-th feature.
            perturbed = sample.flatten().copy()

            if self.dataset.name == "digits" or self.dataset.name == "mnist":
                perturbed[i] = 1 - perturbed[i]
            else:
                perturbed[i] = self.dataset.feature_means[i]

            perturbed = perturbed.reshape(1, *sample.shape)

            # Get the perturbed prediction.
            pred_pert = self.session.run(None, {self.model.graph.input[0].name: perturbed})[0]

            # Compute the sensitivity of the i-th feature.
            sensitivities.append(np.abs(pred_orig - pred_pert).sum())

        # Return the features in ascending order of their sensitivities.
        return np.arange(np.prod(sample.shape))[np.argsort(sensitivities)]

    def __get_smooth_sensitivity(self, sample):
        # Get the sensitivity of each feature.
        sensitivities = self.__get_sensitivity(sample)

        # Reshape to 2D
        sensitivity_grid = sensitivities.reshape(sample.shape[:2])

        # Define 3x3 kernel to average over neighbors (including self)
        kernel = np.array([
            [1, 1, 1],
            [1, 8, 1],
            [1, 1, 1]], dtype=float).reshape((3, 3))
        kernel /= kernel.sum()  # Normalize to average

        # Apply convolution with padding to preserve size
        sensitivity = convolve(sensitivity_grid, kernel, mode='reflect').astype(float)

        # Create a positional weighting mask
        # Example: Gaussian falloff from the center
        account_for_positional_weighting = True

        if account_for_positional_weighting:
            SIGMA = 0.8

            # Create a 2D Gaussian positional weighting mask
            x = np.linspace(-1, 1, sample.shape[0])
            y = np.linspace(-1, 1, sample.shape[1])
            xx, yy = np.meshgrid(x, y)

            # Shape the Gaussian to peak at the center
            positional_weight = np.exp(-(xx**2 + yy**2) / (2 * SIGMA**2))
            sensitivity *= positional_weight

        return sensitivity.flatten()
