"""Macro F0.5 exactly as defined in the problem statement."""
import numpy as np


def entity_scores(pred, truth, beta=0.5):
    """pred, truth: iterables of IDs for ONE Source-1 entity. Returns (f, precision, recall)."""
    p, t = set(pred), set(truth)
    if not t:
        return (1.0, 1.0, 1.0) if not p else (0.0, 0.0, 1.0)
    if not p:
        return 0.0, 0.0, 0.0
    tp = len(p & t)
    prec, rec = tp / len(p), tp / len(t)
    if tp == 0:
        return 0.0, prec, rec
    b2 = beta * beta
    return (1 + b2) * prec * rec / (b2 * prec + rec), prec, rec


def macro_f05(preds, truths, ids):
    """preds / truths: dict id -> list. Every id in `ids` is scored (missing pred = empty)."""
    rows = [entity_scores(preds.get(i, []), truths.get(i, [])) for i in ids]
    a = np.array(rows) if rows else np.zeros((0, 3))
    return {"f05": float(a[:, 0].mean()), "precision": float(a[:, 1].mean()),
            "recall": float(a[:, 2].mean()), "n": len(ids)}, a[:, 0]


def _self_test():
    f, p, r = entity_scores(["S2-00047", "S2-00193", "S3-00812"], ["S2-00047", "S3-00812"])
    assert abs(f - 0.714) < 1e-3, f
    assert entity_scores([], [])[0] == 1.0
    assert entity_scores(["S2-1"], [])[0] == 0.0
    assert entity_scores([], ["S2-1"])[0] == 0.0


_self_test()