"""Measure the phenotype ranking against every published case in Phenopacket Store.

One demo case ranking correctly is an anecdote. This asks the only question that matters for a
phenotype-driven tool: given a real patient's reported terms, where does the gene the authors
actually diagnosed come in a ranking of every gene HPO annotates?

    python -m rarelens_ml.benchmark --phenopackets all_phenopackets.zip

**Read the result with the contamination in mind.** HPO's gene-to-phenotype annotations are
themselves curated from published case reports — quite possibly the very ones being scored here.
The median causal gene already carries every one of its patient's terms, so this measures how well
the ranking retrieves a gene HPO has already been told about. It is an upper bound. A prospective
number, on a patient whose disease gene nobody has annotated yet, would be lower; how much lower
this corpus cannot say.

Ties are the other trap. Scoring by term overlap alone puts many genes on identical scores, so the
honest report is a range: optimistic counts a tie as a win, pessimistic counts every tied gene as
ranked ahead of the right answer. The truth is between them.
"""
import argparse
import json
import statistics
import sys
import zipfile
from collections import defaultdict
from collections.abc import Iterator

from rarelens_ml.hpo import information_content, phenotype_score


def gene_annotations(path: str) -> dict[str, set[str]]:
    """gene -> HPO terms, from the propagated table scripts/load-hpo.py writes (TSV export)."""
    genes: dict[str, set[str]] = defaultdict(set)
    with open(path) as fh:
        for line in fh:
            gene, _, term = line.rstrip("\n").partition("\t")
            if gene and term:
                genes[gene].add(term)
    return genes


def cases(path: str) -> Iterator[tuple[str, list[str]]]:
    """(causal gene, observed HPO terms) for each phenopacket with exactly one causal gene."""
    with zipfile.ZipFile(path) as z:
        for entry in z.namelist():
            if not entry.endswith(".json"):
                continue
            try:
                packet = json.loads(z.read(entry))
            except ValueError:
                continue
            causal = {
                g.get("variantInterpretation", {})
                .get("variationDescriptor", {})
                .get("geneContext", {})
                .get("symbol")
                for i in packet.get("interpretations", [])
                for g in i.get("diagnosis", {}).get("genomicInterpretations", [])
            } - {None}
            terms = [
                f["type"]["id"]
                for f in packet.get("phenotypicFeatures", [])
                if not f.get("excluded")
            ]
            if len(causal) == 1 and terms:
                yield causal.pop(), terms


def rank_of(gene: str, terms: list[str], genes: dict[str, set[str]], ic: dict[str, float],
            default: float) -> tuple[int, int, float]:
    """(optimistic rank, pessimistic rank, the causal gene's own score)."""
    target = phenotype_score(terms, genes[gene], ic, default)
    better = tied = 0
    for other, annotated in genes.items():
        if other == gene:
            continue
        value = phenotype_score(terms, annotated, ic, default)
        if value > target:
            better += 1
        elif value == target:
            tied += 1
    return better + 1, better + tied + 1, target


def report(ranks: list[tuple[int, int, float]], n_genes: int, min_terms: int = 1) -> str:
    n = len(ranks)
    median_score = statistics.median(t for *_, t in ranks)
    lines = [
        (
            f"{n} published cases with at least {min_terms} HPO term(s), ranked against "
            f"{n_genes} genes (random top-1 would be {1 / n_genes:.2%})"
        ),
        (
            f"the causal gene's own phenotype score: median {median_score:.2f}"
            "  <- 1.00 means HPO already carries every one of the patient's terms for that gene"
        ),
    ]
    for label, column in (("optimistic", 0), ("pessimistic", 1)):
        r = [row[column] for row in ranks]
        lines.append(
            f"  {label:12s} top-1 {sum(x == 1 for x in r) / n:6.1%}   "
            f"top-10 {sum(x <= 10 for x in r) / n:6.1%}   "
            f"MRR {sum(1 / x for x in r) / n:.3f}   median rank {statistics.median(r):.0f}"
        )
    return "\n".join(lines)


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--phenopackets", required=True,
                   help="all_phenopackets.zip from a phenopacket-store release")
    p.add_argument("--annotations", required=True,
                   help="TSV of gene<tab>hpo_id, exported from the gene_phenotypes table")
    p.add_argument("--limit", type=int, help="benchmark only the first N cases (a smoke run)")
    p.add_argument("--min-terms", type=int, default=1,
                   help="skip cases with fewer HPO terms; a one-term case can only tie")
    a = p.parse_args()

    genes = gene_annotations(a.annotations)
    ic = information_content(genes)
    default = max(ic.values(), default=1.0)
    ranks = []
    for gene, terms in cases(a.phenopackets):
        if gene in genes and len(terms) >= a.min_terms:
            ranks.append(rank_of(gene, terms, genes, ic, default))
        if a.limit and len(ranks) >= a.limit:
            break
    if not ranks:
        sys.exit("no benchmarkable cases: is --annotations the gene_phenotypes export?")
    print(report(ranks, len(genes), a.min_terms))


if __name__ == "__main__":
    main()