from abc import ABC, abstractmethod
from collections.abc import Callable, Sequence
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    import torch

from visuallm.elements.barchart_element import BarChartElement, PieceInfo
from visuallm.elements.plain_text_element import PlainTextElement
from visuallm.elements.selector_elements import ButtonElement, CheckBoxSubElement


@dataclass
class MetricDescription:
    format: str
    """Format of the value that should be displayed on the page, e.g.
    "{:.4f}" will display only 4 most significant digits of the value
    """
    scalable: bool
    """Whether the bar is also scalable, or should be kept at 100 for the
    metric number to be better visible
    """
    metric_calculation: Callable[[Any, Any], Any]


@dataclass
class GeneratedTextMetric(MetricDescription):
    metric_calculation: Callable[[str, str], Any]
    """Function that given the text generated by the model and the target text
    will generate a value that can be displayed using `self.format`
    """


@dataclass
class ProbsMetric(MetricDescription):
    metric_calculation: Callable[[Any, Any], Any]
    """Function that given the probabilities vectors and target indices will
    generate a value that can be displayed using `self.format`"""


class MetricsMixin(ABC):

    """Add the following two types of elements:

    1. the selectors, which allow to select which metrics should be computed on
    the particular input and expected output pairs
    2. the barchart elements which display the metrics
    """

    def __init__(
        self,
        metrics_on_generated_text: dict[str, GeneratedTextMetric] | None = None,
        metrics_on_probs: dict[str, ProbsMetric] | None = None,
        use_target_metrics: bool = True,
    ):
        """Mixin that adds the following capabilities:
        1. select which metrics to compute
        2. display computed metrics
        3. compute the metrics

        Args:
        ----
            metrics_on_generated_text (Dict[str, GeneratedTextMetric], optional):
                Metrics which are computed on pairs of strings. Defaults to {}.
            metrics_on_probs (Dict[str, ProbsMetric], optional):
                Metrics which are computed on the probabilities of generations. Defaults to {}.
            use_target_metrics (bool): whether to display also the metrics computed on
                targets
        """
        if metrics_on_generated_text is None:
            metrics_on_generated_text = {}
        if metrics_on_probs is None:
            metrics_on_probs = {}
        self._ordering = self.create_ordering(
            metrics_on_generated_text, metrics_on_probs
        )
        self.use_target_metrics = use_target_metrics
        self._prepare_metrics_selection_frontend(self._ordering)
        self._metrics_on_generated_text = metrics_on_generated_text
        self._metrics_on_probs = metrics_on_probs
        self._display_metrics_heading = PlainTextElement(
            content="Metrics on Generated Outputs", is_heading=True
        )
        self._display_metrics_on_predicted_element = BarChartElement(long_contexts=True)
        if self.use_target_metrics:
            self._display_metrics_on_target_heading = PlainTextElement(
                content="Metrics on Target", is_heading=True
            )
            self._display_metrics_on_target_element = BarChartElement(
                long_contexts=True
            )

        self.metric_button_element = ButtonElement(
            subelements=list(self._select_metrics_elements.values()),
            button_text="Select Metrics to Display",
            processing_callback=self.metrics_processing_callback,
        )

    @staticmethod
    def create_ordering(
        metrics_on_generated_text: dict[str, Any], metrics_on_probs: dict[str, Any]
    ):
        """Create ordering of metrics in which the metrics would be displayed on the
        frontend.
        """
        return list(metrics_on_generated_text.keys()) + list(metrics_on_probs.keys())

    def _prepare_metrics_selection_frontend(self, ordering: list[str]):
        """Prepare frontend elements for metrics checkboxes which allow
        the user to select which metrics we should compute for any output
        of the model.
        """
        self._select_metrics_heading = PlainTextElement(
            content="Which Metrics to Display", is_heading=True
        )
        self._select_metrics_elements: dict[str, CheckBoxSubElement] = {}

        for key in ordering:
            self._select_metrics_elements[key] = CheckBoxSubElement(key, True)

    @property
    def metrics_selection_elements(self):
        """Elements which allow the user to select which metrics should be
        displayed.
        """
        if len(self._metrics_on_generated_text) + len(self._metrics_on_probs) > 0:
            return [self._select_metrics_heading, self.metric_button_element]
        else:
            return []

    @property
    def metrics_display_elements(self):
        """Elements that display the metrics on the target and on the
        generated outputs.
        """
        ret = [
            self._display_metrics_heading,
            self._display_metrics_on_predicted_element,
        ]
        if self.use_target_metrics:
            ret = [
                self._display_metrics_on_target_heading,
                self._display_metrics_on_target_element,
            ] + ret
        return ret

    def _compute_n_display_metrics_for_element(
        self,
        generated_text_list: Sequence[str],
        label_text: str,
        probs_encoded_list: Sequence["torch.Tensor"] | Sequence[None],
        generated_encoded_list: Sequence["torch.Tensor"] | Sequence[None],
        element: BarChartElement,
    ):
        """Calculate generation metrics for each element of `generated_text_list` and
        probability metrics for each element of `probs_encoded_list`.

        Args:
        ----
            generated_text_list (Sequence[str]): Sequence with generations of the model.
            label_text (str): Gold output.
            probs_encoded_list (Sequence['torch.Tensor']): Sequence of tensors of shape
            generated_encoded_list (Sequence['torch.Tensor']): Sequence of generated indices of tokens (sequence
                of generated sequences of indices of tokens)
            element (BarChartElement): Element where to display computed metrics.
        """
        piece_infos: list[PieceInfo] = []
        for (
            generated_text,
            probs_encoded,
            generated_ids_encoded,
        ) in zip(
            generated_text_list, probs_encoded_list, generated_encoded_list, strict=True
        ):
            bar_names: list[str] = []
            bar_annotations: list[str] = []
            bar_heights: list[float] = []
            for name in self._ordering:
                if not self._select_metrics_elements[name].value_on_backend:
                    continue
                bar_names.append(name)

                if name in self._metrics_on_generated_text:
                    metric_description: MetricDescription = (
                        self._metrics_on_generated_text[name]
                    )
                    result = metric_description.metric_calculation(
                        generated_text, label_text
                    )
                else:
                    metric_description = self._metrics_on_probs[name]
                    if probs_encoded is None:
                        result = 0
                    else:
                        result = metric_description.metric_calculation(
                            probs_encoded, generated_ids_encoded
                        )

                if metric_description.scalable:
                    bar_heights.append(min(result * 100, 100))
                else:
                    bar_heights.append(100)
                bar_annotations.append(metric_description.format.format(result))

            piece_infos.append(
                PieceInfo(
                    pieceTitle=generated_text,
                    barHeights=bar_heights,
                    barAnnotations=bar_annotations,
                    barNames=bar_names,
                )
            )

        element.set_piece_infos(piece_infos)

    def compute_n_display_metrics_on_predicted(
        self,
        generated_text_list: Sequence[str],
        label_text: str,
        probs_encoded_list: Sequence["torch.Tensor"] | Sequence[None],
        generated_encoded_list: Sequence[None] | Sequence["torch.Tensor"],
    ):
        """Compute metrics on the predictions of the model and
        display them with `self._display_metrics_on_predicted_element`.

        Args:
        ----
            generated_text_list (Sequence[str]): list of generations of the model.
            label_text (str): gold output.
            probs_encoded_list (Sequence['torch.Tensor']): Sequence of tensors with shape (1, sentence_length, vocab_size)
                each depicting the probabilities
            generated_encoded_list (Sequence['torch.Tensor']): Sequence of tensors with shape (sentence_length,) each
                depicting the token ids that vere actually generated
        """
        self._compute_n_display_metrics_for_element(
            generated_text_list,
            label_text,
            probs_encoded_list,
            generated_encoded_list,
            self._display_metrics_on_predicted_element,
        )

    def compute_n_display_metrics_on_target(
        self,
        target: str,
        probs_target: Sequence["torch.Tensor"] | Sequence[None],
        target_encoded: Sequence["torch.Tensor"] | Sequence[None],
    ):
        """Compute metrics on the targets.

        Args:
        ----
            target (str): gold output.
            probs_target (Sequence['torch.Tensor']): Sequence of sequences of probabilities of each
                target token.
            generated_encoded_list (Sequence['torch.Tensor']): Sequence of sequences of ids of each
                target token.
            target_encoded (Sequence['torch.Tensor']): TODO
        """
        if not self.use_target_metrics:
            raise ValueError("This component wasn't configured for target metrics!")
        self._compute_n_display_metrics_for_element(
            [target],
            target,
            probs_target,
            target_encoded,
            self._display_metrics_on_target_element,
        )

    @abstractmethod
    def metrics_processing_callback(self):
        """What to do right after the selection of which metrics to display is
        updated. (This mixin does not automatically calculate metrics as it may
        require e.g. non-trivial manipulation with the inputs. Hence you need to
        implement metrics computation in this method.)
        """
        ...
