"""Micro + per-class precision/recall/F1 accumulators. Port of ner_metrics.py / relation_metrics.py / f1.py.

Per-class tp/fp/fn are accumulated alongside the micro totals so macro-F1 and the rare-class
breakdown can be logged live during training (not only post-hoc in eval.py) -- same information
the v1 exp2 class-weighting runs surfaced to W&B.
"""
from collections import defaultdict
from typing import Dict, Tuple

import torch


def compute_f1(predicted: float, gold: float, matched: float) -> Tuple[float, float, float]:
    precision = matched / predicted if predicted > 0 else 0.0
    recall = matched / gold if gold > 0 else 0.0
    f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0.0
    return precision, recall, f1


class NERMetrics:
    """Exact-match span+label F1, from tensors (no mask: batch_size=1 means no padding)."""

    def __init__(self, num_classes: int, null_label: int = 0):
        self.num_classes = num_classes
        self.null_label = null_label
        self.reset()

    def reset(self):
        self.tp = self.fp = self.fn = 0
        self.per_class = {c: [0, 0, 0] for c in range(self.num_classes) if c != self.null_label}

    def __call__(self, predictions: torch.Tensor, gold_labels: torch.Tensor):
        for c in range(self.num_classes):
            if c == self.null_label:
                continue
            tp = int(((predictions == c) & (gold_labels == c)).sum())
            fp = int(((predictions == c) & (gold_labels != c)).sum())
            fn = int(((predictions != c) & (gold_labels == c)).sum())
            self.tp += tp
            self.fp += fp
            self.fn += fn
            pc = self.per_class[c]
            pc[0] += tp
            pc[1] += fp
            pc[2] += fn

    def get_metric(self, reset: bool = False) -> Tuple[float, float, float]:
        result = compute_f1(self.tp + self.fp, self.tp + self.fn, self.tp)
        if reset:
            self.reset()
        return result

    def get_per_class(self) -> Dict[int, float]:
        """{class index -> F1}. Call BEFORE get_metric(reset=True), which clears the counts."""
        return {c: compute_f1(v[0] + v[1], v[0] + v[2], v[0])[2] for c, v in self.per_class.items()}


class RelationMetrics:
    """Exact-match (span1, span2, label) F1, decoded predictions vs. the full (unfiltered) gold
    relation dict -- see dataset.py's module docstring for why gold relations whose spans exceed
    max_span_width still count as false negatives here."""

    def __init__(self):
        self.reset()

    def reset(self):
        self.total_gold = self.total_predicted = self.total_matched = 0
        self.per_class = defaultdict(lambda: [0, 0, 0])  # label -> [tp, fp, fn]

    def __call__(self, predicted: Dict[Tuple, str], gold: Dict[Tuple, str]):
        self.total_gold += len(gold)
        self.total_predicted += len(predicted)
        for key, label in predicted.items():
            if gold.get(key) == label:
                self.total_matched += 1
                self.per_class[label][0] += 1  # tp
            else:
                self.per_class[label][1] += 1  # fp: label wrong, or key absent from gold
        for key, label in gold.items():
            if predicted.get(key) != label:
                self.per_class[label][2] += 1  # fn

    def get_metric(self, reset: bool = False) -> Tuple[float, float, float]:
        result = compute_f1(self.total_predicted, self.total_gold, self.total_matched)
        if reset:
            self.reset()
        return result

    def get_per_class(self) -> Dict[str, float]:
        """{relation label -> F1}. Call BEFORE get_metric(reset=True), which clears the counts."""
        return {lab: compute_f1(v[0] + v[1], v[0] + v[2], v[0])[2] for lab, v in self.per_class.items()}