#!/usr/bin/env python3
"""Load HPO into the gene_phenotypes and hpo_terms tables: the reference data the ranking uses.

Two source files:

  https://purl.obolibrary.org/obo/hp/hpoa/genes_to_phenotype.txt   (gene -> term, ~20 MB)
  https://purl.obolibrary.org/obo/hp.obo                           (the ontology itself, ~10 MB)

The ontology matters because HPO's gene annotations are direct, and a patient may be described one
level away from whichever term the curator chose. `rarelens_ml.hpo` holds the propagation and
information-content arithmetic, with the tests, since it is what the phenotype half of the ranking
rests on. Run through the ml environment, which is where that package lives: `make hpo`.

Cite the Human Phenotype Ontology when showing results; see docs/data.md.
"""
import argparse
import csv
import io
import os
import sys
import urllib.request

from rarelens_ml.hpo import ancestors_of, information_content, parse_obo, propagate
from sqlalchemy import create_engine, text
from sqlalchemy.engine import make_url

GENES_URL = "https://purl.obolibrary.org/obo/hp/hpoa/genes_to_phenotype.txt"
OBO_URL = "https://purl.obolibrary.org/obo/hp.obo"


def fetch(url: str) -> io.TextIOBase:
    print(f"downloading {url}", file=sys.stderr)
    return io.TextIOWrapper(urllib.request.urlopen(url), encoding="utf-8")


def annotations(handle: io.TextIOBase) -> set[tuple[str, str]]:
    """Unique (gene, term) pairs; the file repeats them once per associated disease."""
    return {
        (row["gene_symbol"][:60], row["hpo_id"][:20])
        for row in csv.DictReader(handle, delimiter="\t")
        if row.get("gene_symbol") and row.get("hpo_id")
    }


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--url", default=GENES_URL)
    p.add_argument("--obo", default=OBO_URL)
    p.add_argument("--file", help="use a local genes_to_phenotype.txt instead of downloading")
    p.add_argument("--obo-file", help="use a local hp.obo instead of downloading")
    a = p.parse_args()

    url = os.environ.get("DATABASE_URL")
    if not url:
        sys.exit("DATABASE_URL is not set")

    with (open(a.obo_file) if a.obo_file else fetch(a.obo)) as handle:
        parents, names = parse_obo(handle)
    ancestors = ancestors_of(parents)
    with (open(a.file) if a.file else fetch(a.url)) as handle:
        direct = annotations(handle)

    genes = propagate(direct, ancestors)
    ic = information_content(genes)
    rows = [(gene, term, names.get(term, term)) for gene, terms in genes.items() for term in terms]
    print(
        f"{len(direct)} direct annotations over {len(genes)} genes -> {len(rows)} after "
        f"propagation; {len(ic)} terms with information content",
        file=sys.stderr,
    )

    engine = create_engine(make_url(url).set(drivername="postgresql+psycopg"))
    with engine.begin() as conn:
        conn.execute(text("TRUNCATE gene_phenotypes RESTART IDENTITY"))
        conn.execute(text("TRUNCATE hpo_terms"))
        cursor = conn.connection.cursor()
        with cursor.copy("COPY gene_phenotypes (gene_symbol, hpo_id, hpo_name) FROM STDIN") as copy:
            for row in rows:
                copy.write_row(row)
        with cursor.copy("COPY hpo_terms (hpo_id, name, ic) FROM STDIN") as copy:
            for term, value in ic.items():
                copy.write_row((term, names.get(term, term), value))
    print(f"loaded {len(rows)} annotations and {len(ic)} terms", file=sys.stderr)


if __name__ == "__main__":
    main()