"""HPO ontology handling: propagation and information content.

Shared by scripts/load-hpo.py, which writes the tables the API ranks against, and
rarelens_ml.benchmark, which scores that ranking. It lives in the package rather than in the
script so the arithmetic underneath the project's main scientific claim is covered by tests.

Two ideas, both standard practice and both absent from the first version of the ranking:

**Propagation.** HPO's gene annotations are direct. A gene linked to "Aortic root aneurysm" is not
also linked to "Aortic aneurysm", so matching case terms by exact ID missed any patient whose
description sat one level away from the curator's chosen term. The annotation propagation rule
says a gene annotated with a term is annotated with all of that term's ancestors; matching then
works in both directions without the ranking knowing the ontology exists.

**Information content.** IC(term) = -ln(share of genes carrying it). After propagation almost
every gene carries "Abnormality of the cardiovascular system", so its IC is near zero, while
"Dilated left subclavian artery" is worth a great deal. Counting terms alike let a patient's
"Global developmental delay" count as much as a near-pathognomonic sign.
"""
import io
import math
from collections import defaultdict
from collections.abc import Iterable

# Terms outside this branch (inheritance, clinical modifiers, frequency) describe how a disease
# behaves rather than what is wrong with the patient, and must not count towards a match.
PHENOTYPIC_ABNORMALITY = "HP:0000118"


def parse_obo(handle: io.TextIOBase) -> tuple[dict[str, set[str]], dict[str, str]]:
    """Each term's direct parents and its name, from hp.obo. Obsolete terms are dropped."""
    parents: dict[str, set[str]] = {}
    names: dict[str, str] = {}
    term_id: str | None = None
    name: str | None = None
    is_a: set[str] = set()
    obsolete = in_term = False

    def flush() -> None:
        if term_id and not obsolete:
            parents[term_id] = is_a
            names[term_id] = name or term_id

    for raw in handle:
        line = raw.rstrip("\n")
        if line.startswith("["):
            flush()
            term_id, name, is_a, obsolete = None, None, set(), False
            in_term = line == "[Term]"
        elif not in_term:
            continue
        elif line.startswith("id: HP:"):
            term_id = line[4:].strip()
        elif line.startswith("name: "):
            name = line[6:].strip()[:200]
        elif line.startswith("is_a: HP:"):
            is_a.add(line[6:].split("!")[0].strip())
        elif line.startswith("is_obsolete: true"):
            obsolete = True
    flush()
    return parents, names


def ancestors_of(parents: dict[str, set[str]]) -> dict[str, set[str]]:
    """Every term's ancestors, itself included.

    Iterative, because HPO is deep enough to exhaust the recursion limit, and in true post-order:
    a term is resolved only once every parent is resolved. A pre-order walk read backwards looks
    like it would do, but on a DAG a term can be visited before one of its parents on another
    branch, and then it silently inherits that parent alone instead of the parent's whole
    lineage. That dropped Camptodactyly and Chiari malformation out of the phenotype branch
    entirely, which is what this shape of bug looks like from the outside.
    """
    cache: dict[str, set[str]] = {}
    for start in parents:
        if start in cache:
            continue
        stack: list[tuple[str, bool]] = [(start, False)]
        while stack:
            node, resolved = stack.pop()
            if node in cache:
                continue
            if resolved:
                found = {node}
                for parent in parents.get(node, ()):
                    found |= cache.get(parent, {parent})  # fallback guards against a cycle
                cache[node] = found
            else:
                stack.append((node, True))
                stack.extend((p, False) for p in parents.get(node, ()) if p not in cache)
    return cache


def propagate(
    direct: Iterable[tuple[str, str]], ancestors: dict[str, set[str]]
) -> dict[str, set[str]]:
    """gene -> its annotated terms plus all their ancestors, within the phenotype branch."""
    genes: dict[str, set[str]] = defaultdict(set)
    for gene, term in direct:
        for node in ancestors.get(term, {term}):
            if node != PHENOTYPIC_ABNORMALITY and PHENOTYPIC_ABNORMALITY in ancestors.get(node, ()):
                genes[gene].add(node)
    return dict(genes)


def information_content(genes: dict[str, set[str]]) -> dict[str, float]:
    """-ln(share of genes carrying the term); 0 for a term every gene has."""
    if not genes:
        return {}
    counts: dict[str, int] = defaultdict(int)
    for terms in genes.values():
        for term in terms:
            counts[term] += 1
    return {term: -math.log(n / len(genes)) for term, n in counts.items()}


def phenotype_score(
    case_terms: Iterable[str], gene_terms: set[str], ic: dict[str, float], default: float
) -> float:
    """Information-content-weighted recall; the same arithmetic as app.services.triage."""
    terms = list(case_terms)
    total = sum(ic.get(t, default) for t in terms)
    if total <= 0:
        return 0.0
    return sum(ic.get(t, default) for t in terms if t in gene_terms) / total