#!/usr/bin/env python3


import warnings
from typing import Callable, Optional, Tuple

import torch
from tqdm import tqdm


def _feature_mask_to_feature_groups_and_counts(
    feature_mask: Tuple[torch.Tensor, ...]
) -> Tuple:
    """This function takes a tuple of feature masks and returns the grouped feature counts and the number of grouped features.

    Args:
        feature_mask (Tuple[torch.Tensor, ...]): Tuple of feature group masks for each feature type.
        For example, for a model with 4 feature types, the input would be a tuple of 4
        tensors where each tensor is of shape (batch_size, seq_length, n_features).
        And the feature type can be anything like token embeddings, position embeddings, bbox embeddings, etc.

    Returns:
        Tuple[List[torch.Tensor], List[float]]:
            The return value is a tuple of two lists. First list gives the grouped feature count which is the total
            number of elements present in each group for each feature type and each sample in the batch.
            The second list gives the total number of unique groups in each feature type and each sample.

            The output is like following:
                grouped_feature_counts = [
                    [ # feature type 1
                        [ [feature_group_1_count, feature_group_2_count, ...] # sample 1
                        [ [feature_group_1_count, feature_group_2_count, ...] # sample 2
                    ]
                ]

                n_grouped_features = [
                    [ # feature type 1
                        # N-feature-groups for sample 1
                        # N-feature-groups for sample 2
                    ]
                ]
    """

    batch_size = feature_mask[0].shape[0]
    grouped_feature_counts = []
    n_grouped_features = []
    for mask in feature_mask:
        grouped_feature_counts.append([])
        n_grouped_features.append([])
        for idx in range(batch_size):
            grouped_feature_counts[-1].append(
                torch.unique_consecutive(mask[idx, :, 0], return_counts=True)[1]
            )
            n_grouped_features[-1].append(len(torch.unique(mask[idx])))
        grouped_feature_counts[-1] = grouped_feature_counts[-1]
        n_grouped_features[-1] = n_grouped_features[-1]
    return grouped_feature_counts, n_grouped_features


def _divide_and_aggregate_metrics(
    inputs: Tuple[torch.Tensor, ...],
    n_perturb_samples: int,
    metric_func: Callable,
    agg_func: Callable = torch.add,
    max_examples_per_batch: int = None,
    show_progress: bool = False,
) -> torch.Tensor:
    r"""
    This function is used to slice large number of samples `n_perturb_samples` per
    input example into smaller pieces, computing the metrics for each small piece and
    aggregating the results across all `n_perturb_samples` per example. The function
    returns overall aggregated metric per sample. The size of each slice is determined
    by the `max_examples_per_batch` input parameter.

    Args:

        inputs (tuple): The original inputs formatted in a tuple that are passed to
                        the metrics function and that are used to compute the
                        attributions for.
        n_perturb_samples (int): The number of samples per example that are used for
                        perturbation purposes for example.
        metric_func (Callable): This function takes the number of samples per
                        input batch and returns an overall metric for each example.
        agg_func (Callable, optional): This function is used to aggregate the
                        metrics across multiple sub-batches and that are
                        generated by `metric_func`.
        max_examples_per_batch (int, optional): The maximum number of allowed examples
                        per batch.
        show_progress (bool, optional): Whether to show the progress bar.
        Returns:

            metric (Tensor): A metric score estimated by `metric_func` per
                        input example.
    """
    bsz = inputs[0].size(0)

    if max_examples_per_batch is not None and (
        max_examples_per_batch // bsz < 1
        or max_examples_per_batch // bsz > n_perturb_samples
    ):
        warnings.warn(
            (
                "`max_examples_per_batch` must be at least equal to the"
                " input batch size and at most to "
                "`input batch size` * `n_perturb_samples`."
                "`max_examples_per_batch` is: {} and the input batch size is: {}."
                "This is necessary because we require that each sub-batch that is used "
                "to compute the metrics, contains at least an instance of "
                "the original example and doesn't exceed the number of "
                "expanded n_perturb_samples."
            ).format(max_examples_per_batch, bsz)
        )

    max_inps_per_batch = (
        n_perturb_samples
        if max_examples_per_batch is None
        else min(max(max_examples_per_batch // bsz, 1), n_perturb_samples)
    )

    current_n_steps = max_inps_per_batch

    metrics_sum = metric_func(max_inps_per_batch, current_n_steps)

    pbar = None
    while current_n_steps < n_perturb_samples:
        if show_progress and pbar is None:
            pbar = tqdm(total=n_perturb_samples, desc="Computing metrics", leave=False)

        current_n_steps += max_inps_per_batch
        if pbar is not None:
            pbar.update(max_inps_per_batch)

        metric = metric_func(
            (
                max_inps_per_batch
                if current_n_steps <= n_perturb_samples
                else max_inps_per_batch - (current_n_steps - n_perturb_samples)
            ),
            min(current_n_steps, n_perturb_samples),
        )

        current_n_steps = min(current_n_steps, n_perturb_samples)

        metrics_sum = agg_func(metrics_sum, metric)
    return metrics_sum


def _divide_and_aggregate_metrics_n_perturbations_per_feature(
    n_perturbations_per_feature: int,
    n_features: int,
    # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters.
    metric_func: Callable,
    # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters.
    agg_func: Callable = torch.add,
    max_features_processed_per_batch: Optional[int] = None,
    show_progress: bool = False,
) -> torch.Tensor:
    r"""
    This function is used to slice large number of samples `n_perturbations_per_feature` per
    input example into smaller pieces, computing the metrics for each small piece and
    aggregating the results all `n_perturbations_per_feature` per example. The function
    returns overall aggregated metric per sample. The size of each slice is determined
    by the `max_features_processed_per_batch` input parameter.
        metric_func (Callable): This function takes the number of samples per
                        input batch and returns an overall metric for each example.
        agg_func (Callable, optional): This function is used to aggregate the
                        metrics across multiple sub-batches and that are
                        generated by `metric_func`.
        max_features_processed_per_batch (int, optional): The maximum number of allowed examples
                        per batch.
        show_progress (bool, optional): Whether to show the

    Returns:

        metric (Tensor): A metric score estimated by `metric_func` per
                    input example.
    """
    if max_features_processed_per_batch is not None and (
        max_features_processed_per_batch // n_perturbations_per_feature < 1
        or max_features_processed_per_batch // n_perturbations_per_feature > n_features
    ):
        warnings.warn(
            (
                "`max_features_processed_per_batch` must be at least equal to the"
                " n_perturbations_per_feature and at most to "
                "`n_perturbations_per_feature` * `n_features`."
                "`max_features_processed_per_batch` is: {} and the input batch size is: {}."
                "This is necessary because we require that each sub-batch that is used "
                "to compute the metrics, contains at least an instance of "
                "the original example and doesn't exceed the number of "
                "expanded n_features."
            ).format(max_features_processed_per_batch, n_perturbations_per_feature)
        )

    max_inps_per_batch = (
        n_features
        if max_features_processed_per_batch is None
        else min(
            max(max_features_processed_per_batch // n_perturbations_per_feature, 1),
            n_features,
        )
    )

    current_n_steps = max_inps_per_batch

    metrics_sum = metric_func(max_inps_per_batch, current_n_steps)

    pbar = None
    while current_n_steps < n_features:
        if show_progress and pbar is None:
            pbar = tqdm(
                total=n_features, desc="Processing features for metric", leave=False
            )

        current_n_steps += max_inps_per_batch
        if pbar is not None:
            pbar.update(max_inps_per_batch)

        metric = metric_func(
            (
                max_inps_per_batch
                if current_n_steps <= n_features
                else max_inps_per_batch - (current_n_steps - n_features)
            ),
            min(current_n_steps, n_features),
        )

        current_n_steps = min(current_n_steps, n_features)

        metrics_sum = agg_func(metrics_sum, metric)
    return metrics_sum


def _divide_and_aggregate_metrics_n_features(
    n_features: int,
    # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters.
    metric_func: Callable,
    # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters.
    agg_func: Callable = torch.add,
    max_features_processed_per_batch: Optional[int] = None,
    show_progress: bool = False,
) -> torch.Tensor:
    r"""
    This function is used to slice large number of samples `n_features` per
    input example into smaller pieces, computing the metrics for each small piece and
    aggregating the results all `n_features` per example. The function
    returns overall aggregated metric per sample. The size of each slice is determined
    by the `max_features_processed_per_batch` input parameter.
        metric_func (Callable): This function takes the number of samples per
                        input batch and returns an overall metric for each example.
        agg_func (Callable, optional): This function is used to aggregate the
                        metrics across multiple sub-batches and that are
                        generated by `metric_func`.
        max_features_processed_per_batch (int, optional): The maximum number of allowed examples
                        per batch.
        show_progress (bool, optional): Whether to show the progress bar.

    Returns:

        metric (Tensor): A metric score estimated by `metric_func` per
                    input example.
    """
    max_inps_per_batch = (
        n_features
        if max_features_processed_per_batch is None
        else min(max_features_processed_per_batch, n_features)
    )

    current_n_steps = max_inps_per_batch

    pbar = None

    metrics_sum = metric_func(max_inps_per_batch, current_n_steps)
    while current_n_steps < n_features:
        if show_progress and pbar is None:
            pbar = tqdm(
                total=n_features, desc="Processing features for metric", leave=False
            )
        current_n_steps += max_inps_per_batch
        if pbar is not None:
            pbar.update(max_inps_per_batch)

        metric = metric_func(
            (
                max_inps_per_batch
                if current_n_steps <= n_features
                else max_inps_per_batch - (current_n_steps - n_features)
            ),
            min(current_n_steps, n_features),
        )

        current_n_steps = min(current_n_steps, n_features)

        metrics_sum = agg_func(metrics_sum, metric)
    return metrics_sum
