from abc import ABC, abstractmethod
from core.explanations.utils import bitmask_to_list
from core.models.BaseModel import BaseModel
from maraboupy import Marabou

import numpy as np


class BaseExplanationValidator(ABC):
    model: BaseModel
    oracle: Marabou.MarabouCore

    label: int

    features = None
    classes = None

    options = Marabou.createOptions(numWorkers=16,
                                    timeoutInSeconds=300,
                                    verbosity=0,
                                    solveWithMILP=False)

    def __init__(self, model, label, shape):
        self.model = model
        self.label = label

        self.__load_marabou_oracle(self.model.path)

        self.features = np.arange(np.prod(shape))
        self.classes = self.oracle.outputVars[0].flatten()

    def is_explanation_valid(self, sample, explanation):
        # Reshape the sample to a 1D array.
        sample = sample.flatten()[:, None]

        # Ensure there are no leftover constraints.
        self.oracle.clearProperty()

        # Ensure the explanation is a set of features.
        if isinstance(explanation, str):
            explanation = set(bitmask_to_list(explanation, self.model.dataset.num_features))

        # Set the input constraints.
        for feature in self.features:
            min_bound, max_bound = self._get_bounds(sample, explanation, feature)
            self.oracle.setLowerBound(feature, max(0, min_bound))
            self.oracle.setUpperBound(feature, min(1, max_bound))

        # The explanation is valid if the prediction cannot change.
        # That is, if fixing the features in 'explanation' and oscillating the rest
        # does not change the classification of the sample.
        is_valid = self._can_prediction_change() == False

        # Clear the constraints.
        self.oracle.clearProperty()

        return is_valid

    def __load_marabou_oracle(self, model_path):
        self.oracle = Marabou.read_onnx(model_path)

        if self.model.last_node.op_type == "Softmax":
            oracle_output = self.model.last_node.input
        else:
            oracle_output = None

        self.oracle = Marabou.read_onnx(filename=model_path, outputNames=oracle_output)

    def _can_prediction_change(self):
        is_model_sat = 'unsat'

        for j, _ in enumerate(self.classes):
            if j == self.label:
                continue

            # The following inequality corresponds to the restriction:
            # output[label] - output[j] <= -1e-6,
            # which is equivalent to:
            # output[j] - output[label] >= 1e-6.
            self.oracle.addInequality([self.classes[self.label], self.classes[j]],
                                            [1, -1], -1e-2,
                                            isProperty=True)

            # If the model is satisfiable, then output[j] can be greater than output[label],
            # so the classification can change to j.
            # If the model is unsatisfiable, then output[j] cannot be greater than output[label],
            # so the classification cannot change to j. To make sure that the classification
            # cannot change, we continue checking the next output class.
            is_model_sat, _, _ = self.oracle.solve(options=self.options, verbose=False)

            # Clear the output constraints.
            self.oracle.additionalEquList.clear()

            if is_model_sat == 'sat' or is_model_sat == 'TIMEOUT':
                break

        return is_model_sat == 'sat'

    @abstractmethod
    def _get_bounds(self, sample, explanation, feature):
        pass
