"""Verbalising graph facts and scoring how well their evidence supports them (NLI).

Papers describe their own work in the first person ("We propose X"), and NLI models do not
treat "This paper proposes X" as entailed by that (measured with
cross-encoder/nli-deberta-v3-small: 0.04 vs 1.00 for "We propose X."). Paper-subject facts are
therefore verbalised in the first person, and an edge's support score is the maximum over its
phrasings. Known weakness: metric names are poorly distinguished ("We report BLEU" entails
"We report ROUGE." at 0.97), so verbatim quote anchoring remains the primary guard.
"""

from __future__ import annotations

from .graph import KG
from .schemas import Edge

TEMPLATES = {
    "proposes": "{s} proposes {o}.",
    "uses": "{s} uses {o}.",
    "evaluates_on": "{s} is evaluated on {o}.",
    "addresses": "{s} addresses the task of {o}.",
    "reports_metric": "{s} reports results in terms of {o}.",
    "compares_with": "{s} is compared with {o}.",
}

PAPER_TEMPLATES = {
    "proposes": ["We propose {o}."],
    "uses": ["We use {o}."],
    "evaluates_on": ["We evaluate on {o}."],
    "addresses": ["This paper addresses the task of {o}.", "We address {o}."],
    "reports_metric": ["We report {o}."],
}


def verbalise(subject_name: str, relation: str, object_name: str, subject_is_paper: bool = False) -> str:
    """Primary phrasing (used in prompts and displays)."""
    return hypotheses(subject_name, relation, object_name, subject_is_paper)[0]


def hypotheses(subject_name: str, relation: str, object_name: str, subject_is_paper: bool = False) -> list[str]:
    if subject_is_paper and relation in PAPER_TEMPLATES:
        return [t.format(o=object_name) for t in PAPER_TEMPLATES[relation]]
    return [TEMPLATES[relation].format(s=subject_name, o=object_name)]


def edge_hypotheses(kg: KG, e: Edge) -> list[str]:
    subj = kg.nodes.get(e.subject)
    return hypotheses(kg.name(e.subject), e.relation, kg.name(e.object), subject_is_paper=bool(subj and subj.type == "Paper"))


def verbalise_edge(kg: KG, e: Edge) -> str:
    return edge_hypotheses(kg, e)[0]


def score_edge_list(kg: KG, edges: list[Edge], nli) -> dict[str, float]:
    """Max entailment probability over each edge's phrasings, given its own evidence quote.

    ``nli`` needs ``score_pairs(list[(premise, hypothesis)]) -> array[n, 2]`` (entail, contradict).
    """
    pairs: list[tuple[str, str]] = []
    owner: list[str] = []
    for e in edges:
        for h in edge_hypotheses(kg, e):
            pairs.append((e.quote, h))
            owner.append(e.edge_id)
    if not pairs:
        return {}
    scores = nli.score_pairs(pairs)
    out: dict[str, float] = {}
    for eid, s in zip(owner, scores):
        out[eid] = max(out.get(eid, 0.0), round(float(s[0]), 4))
    return out


def score_edges(kg: KG, nli) -> dict[str, float]:
    return score_edge_list(kg, list(kg.edges.values()), nli)