from core.explanations.validators.InflatedExplanationValidator import InflatedExplanationValidator
from core.explanations.factories.IndividualExplanationFactory import IndividualExplanationFactory
from core.explanations.utils import bitmask_to_list


class InflatedExplanationFactory(IndividualExplanationFactory):

    def __init__(self, model, label):
        super().__init__(model, label)

    def inflate_explanations(self, explanations, sample, epsilon, eta, ordering):
        inflated_explanations = []

        validator = InflatedExplanationValidator(
            model=self.model,
            label=self.label,
            sample=sample,
            epsilon=epsilon,
            eta=eta
        )

        for explanation in explanations:
            inflated_explanation = InflatedExplanationFactory._inflate_explanation(
                validator, ordering, explanation, sample
            )
            inflated_explanations.append(inflated_explanation)

        return inflated_explanations

    @staticmethod
    def _inflate_explanation(validator, ordering, explanation, sample):
        features = validator.model.traversal_order(sample, ordering)
        sample_flat = sample.flatten()

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

        # Set the bounds given by the explanation.
        bounds = [0] * len(sample)

        for feature in features:
            if feature in explanation:
                lower_bound = sample_flat[feature]
                upper_bound = sample_flat[feature]
            else:
                lower_bound = max(0, sample_flat[feature] - validator.epsilon)
                upper_bound = min(1, sample_flat[feature] + validator.epsilon)
            bounds[feature] = (float(lower_bound), float(upper_bound))

        # Inflate the bounds of the features in the explanation.
        for feature in explanation:
            bounds = InflatedExplanationFactory._expand_sup(validator, bounds, feature, sample_flat)
            bounds = InflatedExplanationFactory._expand_inf(validator, bounds, feature, sample_flat)

        return bounds

    @staticmethod
    def _expand_sup(validator, bounds, feature, sample):
        new_bound = bounds[feature][1]

        while True:
            if new_bound >= 1:
                new_bound = 1
                break

            if new_bound >= sample[feature] + validator.epsilon:
                new_bound = sample[feature] + validator.epsilon
                break

            new_bound += validator.eta
            bounds[feature] = (bounds[feature][0], new_bound)

            if not validator.is_explanation_valid(sample, bounds):
                new_bound -= validator.eta
                break

        bounds[feature] = (bounds[feature][0], new_bound)
        return bounds

    @staticmethod
    def _expand_inf(validator, bounds, feature, sample):
        new_bound = bounds[feature][0]

        while True:
            if new_bound <= 0:
                new_bound = 0
                break

            if new_bound <= sample[feature] - validator.epsilon:
                new_bound = sample[feature] - validator.epsilon
                break

            new_bound -= validator.eta
            bounds[feature] = (new_bound, bounds[feature][1])

            if not validator.is_explanation_valid(sample, bounds):
                new_bound += validator.eta
                break

        bounds[feature] = (new_bound, bounds[feature][1])
        return bounds
