"""Labels the challenge's Gemma extractor reads back in addition to the one a bullet names.

We ran 32 single-diagnosis reports through the extractor (google/gemma-4-31B-it,
temperature 0) and kept the extra labels that came back. Each link was checked
again inside a seven-bullet report; the one that did not hold there was dropped.

A report is therefore read as the closure of its bullets under these links.
When the number of bullets is limited, bullets that imply other predicted labels
are preferred.
"""
import numpy as np

IMPLIES = {
    "chronic_ischemic_stroke":     ["encephalomalacia_gliosis"],
    # lacunar_stroke -> chronic_ischemic_stroke + small_vessel_ischemic_disease held in a
    # single-diagnosis report but only 2 times in 6 inside a seven-bullet report, so it is not used.
    "intraparenchymal_hemorrhage": ["intracranial_hemorrhage"],
    "glioma":                      ["brain_tumor"],
    "brain_metastasis":            ["brain_tumor"],
    "meningioma":                  ["brain_tumor"],
    "intraventricular_tumor":      ["brain_tumor"],
    "pituitary_tumor":             ["brain_tumor"],
    "hydrocephalus_ex_vacuo":      ["cerebral_atrophy", "ventriculomegaly"],
    "craniotomy_craniectomy":      ["postsurgical_changes"],
    "tumor_resection_cavity":      ["postsurgical_changes"],
}


def closure_matrix(names):
    """C[p, c] = 1 if asserting p makes the extractor also report c (transitively)."""
    idx = {n: i for i, n in enumerate(names)}
    L = len(names)
    C = np.eye(L, dtype=bool)
    for p, cs in IMPLIES.items():
        if p not in idx:
            continue
        for c in cs:
            if c in idx:
                C[idx[p], idx[c]] = True
    for _ in range(L):                       # transitive closure (Floyd-Warshall style)
        new = C | (C @ C)
        if (new == C).all():
            break
        C = new
    return C


def apply_closure(asserted, C):
    """asserted: bool[n, L] of what we write; returns what the extractor reads."""
    return (asserted.astype(np.uint8) @ C.astype(np.uint8)) > 0


def select_with_implications(P, thr, cap, C):
    """Greedy choice of at most `cap` bullets, valued over the implication closure.

    A label's value is its margin over its threshold, `P/thr - 1`. A bullet is worth
    the total value of the labels it adds to the closure, so a parent whose children
    are also predicted is chosen first, and a parent that would add unpredicted
    children is penalised by their negative margin.
    """
    n, L = P.shape
    val = P / np.maximum(thr, 1e-9) - 1.0
    Cf = C.astype(np.float64)
    if not cap:
        return apply_closure(val > 0, C)
    have = np.zeros((n, L), dtype=bool)
    picked = np.zeros((n, L), dtype=bool)
    rows = np.arange(n)
    for _ in range(cap):
        gain = (val * ~have) @ Cf.T          # gain[i, d] = value d adds for study i
        gain[picked] = -np.inf
        d = gain.argmax(1)
        g = gain[rows, d]
        take = g > 1e-12
        picked[rows[take], d[take]] = True
        have[take] |= C[d[take]]
        if not take.any():
            break
    return picked