from typing import TypedDict
import datasets
import evaluate
import numpy as np


class Metrics(TypedDict):
    confusion_matrix: np.ndarray
    confusion_matrix_normalized: np.ndarray
    bias: np.ndarray
    prevalence: np.ndarray
    correct: np.ndarray
    class_precision: np.ndarray
    class_recall: np.ndarray
    accuracy: float
    """
    Also known as micro F1.
    """
    macro_recall: float
    macro_precision: float
    macro_precision_prevalence_calibrated: float
    F1: np.ndarray
    macro_F1: float
    macro_F1_2: float
    weighted_F1: float
    chance: np.ndarray
    MCC: float
    """
    Generalized Matthews Correlation Coefficient [_2].
    """
    KAPPA: float
    """
    Cohen's Kappa.
    """

_DESCRIPTION = """
This module provides various metrics for single label classification.
"""

_KWARGS_DESCRIPTION = """
Args:
    predictions (`int`): Predicted label.
    references (`int`): Ground truth label.
    num_classes (`int`): Number of classes.

Returns:
    metric_dict (`Metrics`)
"""

_CITATION = """
@article{10.1162/tacl_a_00675,
    author = {Opitz, Juri},
    title = {A Closer Look at Classification Evaluation Metrics and a Critical
                    Reflection of Common Evaluation Practice},
    journal = {Transactions of the Association for Computational Linguistics},
    volume = {12},
    pages = {820-836},
    year = {2024},
    month = {06},
    issn = {2307-387X},
    doi = {10.1162/tacl_a_00675},
}
@article{GORODKIN2004367,
    title = {Comparing two K-category assignments by a K-category correlation coefficient},
    journal = {Computational Biology and Chemistry},
    volume = {28},
    number = {5},
    pages = {367-374},
    year = {2004},
    issn = {1476-9271},
    doi = {https://doi.org/10.1016/j.compbiolchem.2004.09.006},
    url = {https://www.sciencedirect.com/science/article/pii/S1476927104000799},
    author = {J. Gorodkin},
}
"""

@evaluate.utils.file_utils.add_start_docstrings(_DESCRIPTION, _KWARGS_DESCRIPTION)
class EvalSingleClassification(evaluate.Metric):
    def _info(self):
        return evaluate.MetricInfo(
            description = _DESCRIPTION,
            citation = _CITATION,
            inputs_description = _KWARGS_DESCRIPTION,
            features = datasets.Features(
                {
                    "predictions": datasets.Value("int32"),
                    "references": datasets.Value("int32"),
                }
            )
        )

    def _compute(self, predictions, references, num_classes: int):
        confusion_matrix = np.zeros((num_classes, num_classes), dtype=int)
        # Shape: (n_classes_pred, n_classes_label)
        for pred, label in zip(predictions, references):
            confusion_matrix[pred, label] += 1

        confusion_matrix_normalized = confusion_matrix / confusion_matrix.sum()

        bias = np.sum(confusion_matrix, axis=1)
        prevalence = np.sum(confusion_matrix, axis=0)
        prevalence_calibration = confusion_matrix.sum() / prevalence / num_classes
        correct = np.diag(confusion_matrix)
        class_precision = correct / bias
        class_recall = correct / prevalence
        accuracy = correct.sum() / confusion_matrix.sum()
        macro_recall = class_recall.mean()
        macro_precision = class_precision.mean()
        macro_precision_prevalence_calibrated = (
            (prevalence_calibration * correct)
            / (confusion_matrix * prevalence_calibration[None, :]).sum(axis=1)
        ) / num_classes
        F1 = correct / (bias + prevalence + EPS)
        macro_F1 = 2 * F1.sum() / num_classes
        macro_F1_2 = (
            2 * (macro_precision * macro_recall) / (macro_precision + macro_recall + EPS)
        )
        weighted_F1 = (prevalence * F1).sum() / (prevalence.sum() + EPS)

        chance = prevalence.T @ bias

        MCC = (accuracy - chance) / np.sqrt(
            (1 - bias.T @ bias) * (1 - prevalence.T @ prevalence)
        )
        KAPPA = (accuracy - chance) / (1 - chance)

        return {
            "confusion_matrix": confusion_matrix,
            "confusion_matrix_normalized": confusion_matrix_normalized,
            "bias": bias,
            "prevalence": prevalence,
            "correct": correct,
            "class_precision": class_precision,
            "class_recall": class_recall,
            "accuracy": accuracy,
            "macro_recall": macro_recall,
            "macro_precision": macro_precision,
            "macro_precision_prevalence_calibrated": macro_precision_prevalence_calibrated,
            "F1": F1,
            "macro_F1": macro_F1,
            "macro_F1_2": macro_F1_2,
            "weighted_F1": weighted_F1,
            "chance": chance,
            "MCC": MCC,
            "KAPPA": KAPPA,
        }