"""The ontology arithmetic sits underneath the phenotype half of the ranking, so it is tested."""
import io
import math
import random

import pytest

from rarelens_ml.hpo import (
    PHENOTYPIC_ABNORMALITY,
    ancestors_of,
    information_content,
    parse_obo,
    phenotype_score,
    propagate,
)

OBO = f"""format-version: 1.2

[Term]
id: {PHENOTYPIC_ABNORMALITY}
name: Phenotypic abnormality

[Term]
id: HP:0001
name: Abnormality of the vasculature
is_a: {PHENOTYPIC_ABNORMALITY} ! Phenotypic abnormality

[Term]
id: HP:0002
name: Aortic aneurysm
is_a: HP:0001 ! Abnormality of the vasculature

[Term]
id: HP:0003
name: Aortic root aneurysm
is_a: HP:0002 ! Aortic aneurysm

[Term]
id: HP:0004
name: Autosomal dominant inheritance

[Term]
id: HP:0005
name: Obsolete thing
is_a: HP:0001 ! Abnormality of the vasculature
is_obsolete: true
"""


def ontology() -> tuple[dict[str, set[str]], dict[str, str]]:
    return parse_obo(io.StringIO(OBO))


def test_parse_obo_reads_parents_and_drops_obsolete_terms() -> None:
    parents, names = ontology()
    assert parents["HP:0003"] == {"HP:0002"}
    assert names["HP:0002"] == "Aortic aneurysm"
    assert "HP:0005" not in parents


def test_ancestors_include_the_term_itself_and_the_whole_lineage() -> None:
    ancestors = ancestors_of(ontology()[0])
    assert ancestors["HP:0003"] == {"HP:0003", "HP:0002", "HP:0001", PHENOTYPIC_ABNORMALITY}
    assert ancestors["HP:0004"] == {"HP:0004"}  # its own branch, not under phenotypic abnormality


def closure(parents: dict[str, set[str]]) -> dict[str, set[str]]:
    """Reference transitive closure by relaxation: obviously correct, too slow for 20k terms."""
    result = {node: {node} | set(ps) for node, ps in parents.items()}
    changed = True
    while changed:
        changed = False
        for node, found in result.items():
            grown = set(found)
            for parent in found - {node}:
                grown |= result.get(parent, {parent})
            if grown != found:
                result[node] = grown
                changed = True
    return result


def test_ancestors_match_a_reference_closure_on_a_tangled_dag() -> None:
    """The regression this guards cost 399 HPO terms, Camptodactyly and Chiari malformation among
    them: on a DAG a term can be reached before one of its parents, and the old walk then gave it
    that parent alone instead of the parent's whole lineage. It only shows up when a node shares
    ancestors by several routes, so the test needs a genuinely tangled graph rather than a
    hand-drawn diamond.
    """
    rng = random.Random(0)
    nodes = [PHENOTYPIC_ABNORMALITY] + [f"HP:{i:04d}" for i in range(1, 80)]
    parents = {PHENOTYPIC_ABNORMALITY: set()}
    for i, node in enumerate(nodes[1:], start=1):
        # only earlier nodes may be parents, which keeps it acyclic
        parents[node] = set(rng.sample(nodes[:i], k=min(i, rng.randint(1, 3))))

    for _ in range(5):  # dict order decides the traversal, so try several
        shuffled = list(parents.items())
        rng.shuffle(shuffled)
        assert ancestors_of(dict(shuffled)) == closure(dict(shuffled))


def test_every_descendant_of_the_root_keeps_the_root() -> None:
    """The property that actually matters: losing it drops the term out of the phenotype branch."""
    parents = {
        PHENOTYPIC_ABNORMALITY: set(),
        "HP:P": {PHENOTYPIC_ABNORMALITY},
        "HP:X": {"HP:P"},
        "HP:N": {"HP:P"},
        "HP:A": {"HP:X", "HP:N"},
    }
    ancestors = ancestors_of(parents)
    for term in ("HP:P", "HP:X", "HP:N", "HP:A"):
        assert PHENOTYPIC_ABNORMALITY in ancestors[term], term


def test_propagation_lets_a_parent_term_match_a_gene_annotated_with_a_child() -> None:
    ancestors = ancestors_of(ontology()[0])
    genes = propagate([("TGFBR2", "HP:0003")], ancestors)
    assert genes["TGFBR2"] == {"HP:0003", "HP:0002", "HP:0001"}


def test_propagation_drops_the_root_and_anything_outside_the_phenotype_branch() -> None:
    ancestors = ancestors_of(ontology()[0])
    genes = propagate([("A", "HP:0003"), ("A", "HP:0004")], ancestors)
    assert PHENOTYPIC_ABNORMALITY not in genes["A"]  # every gene has it; it carries no information
    assert "HP:0004" not in genes["A"]               # inheritance is not a patient finding


def test_information_content_is_zero_for_a_term_every_gene_carries() -> None:
    ic = information_content({"A": {"HP:1", "HP:2"}, "B": {"HP:1"}, "C": {"HP:1"}})
    assert ic["HP:1"] == pytest.approx(0.0)
    assert ic["HP:2"] == pytest.approx(math.log(3))


def test_phenotype_score_weights_by_specificity() -> None:
    ic = {"HP:common": 0.1, "HP:rare": 6.0}
    terms = ["HP:common", "HP:rare"]
    assert phenotype_score(terms, {"HP:rare"}, ic, 1.0) == pytest.approx(6.0 / 6.1)
    assert phenotype_score(terms, {"HP:common"}, ic, 1.0) == pytest.approx(0.1 / 6.1)


def test_phenotype_score_treats_an_unscored_term_as_maximally_specific() -> None:
    """It can never match, so it must depress every gene equally rather than vanish."""
    assert phenotype_score(["HP:1", "HP:unknown"], {"HP:1"}, {"HP:1": 5.0}, 5.0) == pytest.approx(0.5)