from typing import Optional

import torch
from torchvision import transforms
from scipy.stats import spearmanr

from pnpxai.core._types import Model
from pnpxai.explainers.base import Explainer
from pnpxai.evaluator.metrics.base import Metric


class MuFidelity(Metric):
    """
    Computes the MuFidelity metric for attributions.

    Given a `model` and `inputs`, mufidelity of `model` to an explainer at `inputs` is calculated by
    a correlation between difference of predictions and attributions of maked inputs,
    ``evaluation = corr(pred_diff, masked_attr)``.

    The masked inputs are generated by masking `subset_mask` to noised `inputs`,
    ``masked = perturbed * subset_mask + (1.0 - subset_mask) * baseline``

    Args:
        model (Model): The model to evaluate.
        explainer (Optional[Explainer]): The explainer to evaluate.
        n_perturb (int): Number of perturbations to generate.
        noise_scale (int): Scale factor for Gaussian random noise.
        batch_size (int): Batch size for model evaluation.
        grid_size (int): Size of the grid for creating subsets.
        baseline (Union[float, torch.Tensor]): Baseline value for masked subsets.
        mask_agg_dim (Optional[int]): Dimension to aggregate masks.
        **kwargs: Additional kwargs to compute metric in an evaluator. Not required for single usage.

    Reference:
        U. Bhatt, A. Weller, and J. M. F. Moura. Evaluating and aggregating feature-based model attributions. In Proceedings of the IJCAI (2020).
    """

    def __init__(
        self,
        model: Model,
        explainer: Optional[Explainer] = None,
        n_perturb: int = 150,
        noise_scale: float = 0.2,
        batch_size: int = 32,
        grid_size: int = 9,
        baseline: float = 0.,
        mask_agg_dim: Optional[int] = None,
    ):
        super().__init__(model, explainer)
        self.n_perturb = n_perturb
        self.noise_scale = noise_scale
        self.batch_size = batch_size
        self.grid_size = grid_size
        self.baseline = baseline
        self.mask_agg_dim = mask_agg_dim

    def evaluate(
        self,
        inputs: torch.Tensor,
        targets: torch.Tensor,
        attributions: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """
        Args:
            inputs (torch.Tensor): The input data (N x C x H x W).
            targets (torch.Tensor): The target labels for the inputs (N x 1).
            attributions (torch.Tensor): The attributions of the inputs.

        Returns:
            torch.Tensor: The result of the metric evaluation.
        """
        attributions = self._get_attributions(inputs, targets, attributions).to(self.device)

        outputs = self.model(inputs)
        n_classes = outputs.shape[-1]
        predictions = (
            outputs.to(self.device) * torch.eye(n_classes).to(self.device)[targets.to(self.device)]
        ).sum(dim=-1).detach()

        # input, target, attr, pred
        evaluations = []
        zipped = zip(inputs, targets, attributions, predictions)
        for input, target, attr, pred in zipped:
            repeated = torch.stack([input]*self.n_perturb).to(self.device)
            # Add Gaussian random noise
            std, mean = torch.std_mean(repeated)
            noise = torch.randn_like(repeated).to(self.device) * std.to(self.device) + mean.to(self.device)
            perturbed = self.noise_scale * noise + repeated
            perturbed = torch.minimum(repeated, perturbed)
            perturbed = torch.maximum(repeated-1, perturbed)

            # prepare the random masks that will designate the modified subset (S in original equation)
            subset_size = int(self.grid_size ** 2 * self.noise_scale)
            subset_mask = torch.randn(
                (self.n_perturb, self.grid_size ** 2)).to(self.device)
            subset_mask = torch.argsort(subset_mask, dim=-1) > subset_size
            subset_mask = torch.reshape(subset_mask.type(
                torch.float32), (self.n_perturb, 1, self.grid_size, self.grid_size))
            subset_mask = transforms.Resize(
                perturbed.shape[-2:],
                transforms.InterpolationMode("nearest")
            ).forward(subset_mask)
            if self.mask_agg_dim is not None:
                subset_mask = subset_mask.mean(dim=self.mask_agg_dim)

            # Use the masks to set the selected subsets to baseline state
            masked = perturbed * subset_mask + \
                (1.0 - subset_mask) * self.baseline
            masked = masked.to(self.device)

            masked_output = _forward_batch(self.model, masked, self.batch_size)
            pred_diff = pred - masked_output[:, target]

            masked_attr = (attr * (1.0 - subset_mask))\
                .sum(dim=tuple(range(1, subset_mask.ndim)))

            corr, _ = spearmanr(
                pred_diff.cpu().detach().numpy(),
                masked_attr.cpu().detach().numpy(),
            )
            evaluations.append(corr)
        return torch.tensor(evaluations)


def _forward_batch(model, inputs, batch_size) -> torch.Tensor:
    training_mode = model.training
    model.eval()
    device = next(model.parameters()).device
    outputs = []
    indices = list(range(len(inputs)))
    while indices:
        curr_indices, indices = indices[:batch_size], indices[batch_size:]
        outputs.append(model(inputs[curr_indices].to(device)))
    model.training = training_mode
    return torch.cat(outputs).detach()

