from typing import Optional

import evaluate


class DirectComputeMetric(evaluate.Metric):
    """
    Base class for metrics that directly compute the score from the predictions and references without add_batch
    """

    def compute(self, *, predictions=None, references=None, **kwargs) -> Optional[dict]:
        """Compute the evaluation module.

        Usage of positional arguments is not allowed to prevent mistakes.

        Args:
            predictions (`list/array/tensor`, *optional*):
                Predictions.
            references (`list/array/tensor`, *optional*):
                References.
            **kwargs (optional):
                Keyword arguments that will be forwarded to the evaluation module [`~evaluate.EvaluationModule.compute`]
                method (see details in the docstring).

        Return:
            `dict` or `None`

            - Dictionary with the results if this evaluation module is run on the main process (`process_id == 0`).
            - `None` if the evaluation module is not run on the main process (`process_id != 0`).

        ```py
        >>> import evaluate
        >>> accuracy =  evaluate.load("accuracy")
        >>> accuracy.compute(predictions=[0, 1, 1, 0], references=[0, 1, 0, 1])
        ```
        """
        all_kwargs = {"predictions": predictions, "references": references, **kwargs}
        if predictions is None and references is None:
            missing_kwargs = {k: None for k in self._feature_names() if k not in all_kwargs}
            all_kwargs.update(missing_kwargs)
        else:
            missing_inputs = [k for k in self._feature_names() if k not in all_kwargs]
            if missing_inputs:
                raise ValueError(
                    f"Evaluation module inputs are missing: {missing_inputs}. All required inputs are {list(self._feature_names())}"
                )
        inputs = {input_name: all_kwargs[input_name] for input_name in self._feature_names()}
        compute_kwargs = {k: kwargs[k] for k in kwargs if k not in self._feature_names()}
        return self._compute(**inputs, **compute_kwargs)