from core.explanations.analyzers.BaseExplanationAnalyzer import BaseExplanationAnalyzer
from core.explanations.utils import (
    bitmask_to_list,
    calc_samples_prototype, calc_samples_std
)

import numpy as np


class ClusterExplanationAnalyzer(BaseExplanationAnalyzer):

    @staticmethod
    def calc_explanation_coverage(explanation, samples, prototype, epsilon, delta):
        if len(samples) == 0:
            return 0

        # Convert samples to numpy array and flatten
        samples_array = np.array(samples)
        original_shape = samples_array.shape
        samples_flat = samples_array.reshape(original_shape[0], -1)  # (n_samples, n_features)

        sample_size = samples_flat.shape[1]

        # Handle string explanation conversion
        if isinstance(explanation, str):
            explanation = bitmask_to_list(explanation, sample_size)

        # Convert explanation to boolean mask for vectorized operations
        explanation_mask = np.zeros(sample_size, dtype=bool)
        explanation_mask[explanation] = True

        # Compute statistics once
        samples_prot = calc_samples_prototype(samples, prototype)
        samples_std = calc_samples_std(samples)

        # Create bounds arrays
        # For features NOT in explanation: use epsilon bounds
        epsilon_bounds = np.column_stack([
            samples_prot - epsilon,
            samples_prot + epsilon
        ])

        # For features IN explanation: use gamma bounds
        gamma = delta * epsilon

        if len(samples) > 1:
            gamma *= samples_std

        gamma_bounds = np.column_stack([
            samples_prot - gamma,
            samples_prot + gamma
        ])

        # Combine bounds based on explanation mask
        lower_bounds = np.where(explanation_mask, gamma_bounds[:, 0], epsilon_bounds[:, 0])
        upper_bounds = np.where(explanation_mask, gamma_bounds[:, 1], epsilon_bounds[:, 1])

        # Vectorized bounds checking for all samples at once
        within_bounds = (samples_flat >= lower_bounds) & (samples_flat <= upper_bounds)

        # Calculate feature probabilities for each sample
        feature_probs = np.mean(within_bounds, axis=1)  # Mean across features for each sample

        # Average across all samples
        coverage = np.mean(feature_probs)

        return coverage

    @staticmethod
    def calc_explanation_robustness(explanation, samples, prototype, maximal_bounds, epsilon, delta):
        sample_size = len(samples[0].flatten())

        if isinstance(explanation, str):
            explanation = bitmask_to_list(explanation, sample_size)

        volume = 0
        max_volume = 0

        bounds = ClusterExplanationAnalyzer.calc_explanation_bounds(explanation, samples, prototype, epsilon, delta)

        for feature in explanation:
            max_lower_bound, max_upper_bound = maximal_bounds[feature]
            max_volume += max_upper_bound - max_lower_bound

            lower_bound, upper_bound = bounds[feature]
            volume += upper_bound - lower_bound

        if max_volume == 0:
            return 0

        return volume / max_volume

    @staticmethod
    def calc_explanation_bounds(explanation, samples, prototype, epsilon, delta):
        sample_size = len(samples[0].flatten())

        if isinstance(explanation, str):
            explanation = bitmask_to_list(explanation, sample_size)

        samples_prot = calc_samples_prototype(samples, prototype)
        samples_std = calc_samples_std(samples).flatten()

        bounds = []

        for feature in range(sample_size):
            if feature in explanation:
                gamma = delta * epsilon

                if len(samples) > 1:
                    gamma *= samples_std[feature]

                lower_bound = max(0, samples_prot[feature] - gamma)
                upper_bound = min(1, samples_prot[feature] + gamma)
                bounds.append((lower_bound, upper_bound))
            else:
                lower_bound = max(0, samples_prot[feature] - epsilon)
                upper_bound = min(1, samples_prot[feature] + epsilon)
                bounds.append((lower_bound, upper_bound))

        return bounds
