#!/usr/bin/env python3
"""Extract GO's OWN curated TCDB cross-references from the GO ontology.

Unlike the noisy TCDB-published ``go.py`` dump, GO curators have attached TC
(Transporter Classification) references directly to GO **molecular-function**
terms, in two forms:

  * a term-level ``xref: TC:<id>`` clause  (the strong "this GO term corresponds
    to this TC entry" assertion);
  * a definition dbxref ``def: "..." [TC:<id>]``  (the TC entry cited as the
    basis / source of the term's textual definition).

These xrefs are largely **neglected**, and the term-xref vs definition-dbxref
distinction is not meaningful -- so this script treats both alike and emits them
as unreviewed **sources / leads** (``skos:relatedMatch`` +
``semapv:UnspecifiedMatching``), NOT as an asserted mapping. Whether each lead is
safe to **propagate** (a protein carrying the TC number inheriting the GO term)
must be curated by hand, per entry; those reviewed judgments live in the sibling
``tc2go.sssom.yaml``. This script only surfaces the candidates.

Every row is a GO curator's xref; nothing is hand-typed here.

Usage::

    uv run python extract_go_tc_xrefs.py --stats
    uv run python extract_go_tc_xrefs.py -o tc2go.from_go.sssom.yaml
"""
from __future__ import annotations

import argparse
import re
from pathlib import Path

import urllib.request

GO_OBO_URL = "http://purl.obolibrary.org/obo/go/go-basic.obo"
DATA = Path(__file__).resolve().parent / "data"
CACHE_OBO = DATA / "go-basic.obo"
TC_RE = re.compile(r"TC:[0-9][0-9A-Za-z.\-]*")


def load_obo(refresh: bool = False) -> str:
    if refresh or not CACHE_OBO.exists():
        DATA.mkdir(exist_ok=True)
        req = urllib.request.Request(GO_OBO_URL, headers={"User-Agent": "ai-gene-review-tcdb"})
        with urllib.request.urlopen(req, timeout=180) as r:  # noqa: S310
            CACHE_OBO.write_bytes(r.read())
        print(f"# fetched go-basic.obo -> {CACHE_OBO}")
    return CACHE_OBO.read_text()


def _stanzas(obo: str) -> list[list[str]]:
    """Split the OBO text into per-stanza line lists (each starting after a header line)."""
    blocks: list[list[str]] = []
    cur: list[str] | None = None
    for line in obo.splitlines():
        if line.startswith("[") and line.endswith("]"):
            if cur is not None:
                blocks.append(cur)
            cur = [line]
        elif cur is not None:
            cur.append(line)
    if cur is not None:
        blocks.append(cur)
    return blocks


def extract(obo: str) -> list[dict]:
    """Return [{tc, go, go_label, aspect, source}] for every NON-OBSOLETE GO-term TC reference.

    Obsolete terms are skipped: they retain their old TC xrefs but are detached from the
    molecular-function hierarchy, so they are not valid mapping targets.
    """
    out: list[dict] = []
    for block in _stanzas(obo):
        if not block or block[0] != "[Term]":
            continue
        go = name = ns = None
        obsolete = False
        refs: list[tuple[str, str]] = []  # (tc, source)
        for line in block[1:]:
            if line.startswith("id: GO:"):
                go = line[4:].strip()
            elif line.startswith("name:"):
                name = line.split(":", 1)[1].strip()
            elif line.startswith("namespace:"):
                ns = line.split(":", 1)[1].strip()
            elif line.startswith("is_obsolete: true"):
                obsolete = True
            elif line.startswith("def:"):
                refs += [(tc[3:], "definition") for tc in TC_RE.findall(line)]
            elif line.startswith("xref: TC:"):
                refs.append((line.split("xref:", 1)[1].strip().split()[0][3:], "term_xref"))
        if not go or obsolete:
            continue
        for tc, source in refs:
            out.append({"tc": tc, "go": go, "go_label": name, "aspect": ns, "source": source})
    return out


def cmd_stats(rows: list[dict]) -> None:
    import collections
    go_terms = {r["go"] for r in rows}
    fams = {".".join(r["tc"].split(".")[:3]) for r in rows}
    systems = {r["tc"] for r in rows}
    by_src = collections.Counter(r["source"] for r in rows)
    aspect_by_term = {r["go"]: r["aspect"] for r in rows}
    by_aspect = collections.Counter(aspect_by_term[g] for g in go_terms)
    print("== GO's curated TCDB cross-references ==")
    print(f"total TC references : {len(rows)}  ({by_src['term_xref']} term-level xref:, "
          f"{by_src['definition']} definition dbxref)")
    print(f"distinct GO terms   : {len(go_terms)}")
    print(f"distinct TC systems : {len(systems)}")
    print(f"distinct TC families: {len(fams)}")
    print("GO term aspects     : " + ", ".join(f"{k}={v}" for k, v in by_aspect.items()))


HEADER = """# TC(system) -> GO(molecular function) SOURCES (SSSOM) -- GENERATED from GO's TC xrefs
#
# Auto-generated by projects/TCDB/extract_go_tc_xrefs.py from the live GO ontology (go-basic.obo).
# Each row is a TC reference GO carries on a molecular-function term, inverted to TC->GO. These are
# LEADS / SOURCES, not an asserted mapping: GO's TC xrefs are largely neglected, and the term-level
# `xref: TC:` vs definition-dbxref (`def "..." [TC:]`) distinction is NOT meaningful -- both are
# treated alike here. Every row is therefore skos:relatedMatch + semapv:UnspecifiedMatching, meaning
# "a curator once linked these; whether a protein carrying this TC should INHERIT this GO term
# (propagation) has NOT been established". Propagation must be curated by hand, per entry -- see the
# review-backed judgments in the sibling tc2go.sssom.yaml. No row here is hand-typed.
"""


def to_sssom(rows: list[dict]) -> dict:
    seen = set()
    mappings = []
    for r in sorted(rows, key=lambda r: (r["tc"], r["go"])):
        key = (r["tc"], r["go"])  # term-vs-def distinction dropped (not meaningful)
        if key in seen:
            continue
        seen.add(key)
        mappings.append({
            "subject_id": f"TC:{r['tc']}",
            "subject_label": "",
            "predicate_id": "skos:relatedMatch",
            "predicate_label": "related match",
            "object_id": r["go"],
            "object_label": r["go_label"],
            "mapping_justification": "semapv:UnspecifiedMatching",
            "comment": (
                "GO carries a TC cross-reference on this term (source: go-basic.obo). Unreviewed "
                "lead -- propagation to proteins NOT yet curated (see tc2go.sssom.yaml)."
            ),
        })
    return {
        "curie_map": {
            "TC": "http://www.tcdb.org/search/result.php?tc=",
            "GO": "http://purl.obolibrary.org/obo/GO_",
            "skos": "http://www.w3.org/2004/02/skos/core#",
            "semapv": "https://w3id.org/semapv/vocab/",
        },
        "mapping_set_id": "https://w3id.org/ai4curation/ai-gene-review/mappings/tc2go-from-go",
        "mapping_set_title": "TC -> GO molecular function SOURCES (GO's neglected TC cross-references)",
        "mapping_set_description": (
            "GO's TCDB cross-references, extracted from go-basic.obo and inverted to TC(system) -> "
            "GO(molecular function). These are unreviewed SOURCES/LEADS, not an asserted mapping: "
            "GO's TC xrefs are largely neglected and the term-xref vs definition-dbxref distinction "
            "is not meaningful, so all rows are skos:relatedMatch + semapv:UnspecifiedMatching. "
            "Whether each is safe to PROPAGATE (a protein with this TC inheriting this GO term) must "
            "be curated by hand -- the reviewed judgments live in tc2go.sssom.yaml."
        ),
        "license": "https://creativecommons.org/licenses/by/4.0/",
        "creator_label": ["Gene Ontology curators (extracted by extract_go_tc_xrefs.py)"],
        "subject_source": "tcdb",
        "object_source": "GO",
        "mappings": mappings,
    }


def main(argv: list[str] | None = None) -> int:
    import yaml

    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--stats", action="store_true", help="print counts only")
    ap.add_argument("-o", "--output", type=Path, help="write SSSOM to this path")
    ap.add_argument("--refresh", action="store_true", help="re-fetch go-basic.obo")
    args = ap.parse_args(argv)
    rows = extract(load_obo(refresh=args.refresh))
    if args.stats or not args.output:
        cmd_stats(rows)
    if args.output:
        doc = to_sssom(rows)
        body = yaml.safe_dump(doc, sort_keys=False, default_flow_style=False, width=100, allow_unicode=True)
        args.output.write_text(HEADER + body)
        print(f"# wrote {len(doc['mappings'])} rows -> {args.output}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
