#!/usr/bin/env python3
"""Generate a candidate TC-family -> GO(molecular function) mapping from TCDB's go.py.

There is no public ``tc2go`` external2go mapping. TCDB *does* publish per-protein
GO associations (``go.py``), but that dump is a noisy, multi-aspect aggregation of
member-protein annotations. This script distils the **clean, in-scope slice**:

  1. keep only rows whose GO term is a *transporter-activity molecular function*
     (is_a closure of GO:0022857 / GO:0005215, fetched live from QuickGO);
  2. drop obsolete GO terms and attach the current QuickGO label;
  3. aggregate the surviving 5-level TC numbers up to the **TC family** (3-level),
     counting how many distinct members support each (family, GO) pair;
  4. emit SSSOM: a family associated with a single in-scope MF term -> ``skos:exactMatch``
     (mono-specific: one GO term fits the whole family); a family with several -> ``skos:narrowMatch``
     (poly-specific; the GO term applies to a subfamily). Mono-specificity is judged **before**
     the ``--min-support`` filter, so a poly-specific family is never promoted to ``exactMatch``
     just because its other substrate terms are thinly supported.

Mirrors projects/GLYCOBIOLOGY/build_cazy2go.py. **Machine-derived: every row is a
live join of TCDB's own assertion with the QuickGO ontology -- no row is hand-typed.**
The hand-curated, review-backed seed lives in the sibling ``tc2go.sssom.yaml``.

Usage::

    uv run python build_tc2go.py -o tc2go.generated.sssom.yaml [--min-support 1]
"""
from __future__ import annotations

import argparse
import collections
import json
import urllib.request
from pathlib import Path

from tcdb_go_probe import load_tcdb_go, transporter_mf_closure  # reuse the probe helpers

QUICKGO_TERMS = "https://www.ebi.ac.uk/QuickGO/services/ontology/go/terms/{ids}"
HERE = Path(__file__).resolve().parent


def go_labels(ids: list[str]) -> dict[str, dict]:
    """Fetch {id: {name, aspect, obsolete}} from QuickGO in batches of 400."""
    out: dict[str, dict] = {}
    for i in range(0, len(ids), 400):
        chunk = ids[i : i + 400]
        url = QUICKGO_TERMS.format(ids=",".join(chunk))
        req = urllib.request.Request(url, headers={"Accept": "application/json"})
        with urllib.request.urlopen(req, timeout=90) as r:  # noqa: S310
            data = json.loads(r.read().decode("utf-8", "replace"))
        for res in data.get("results", []):
            out[res["id"]] = {
                "name": res.get("name", ""),
                "aspect": res.get("aspect", ""),
                "obsolete": bool(res.get("isObsolete", False)),
            }
    return out


def build(min_support: int) -> list[dict]:
    rows = load_tcdb_go()
    mf = transporter_mf_closure()
    # family -> Counter(go_id -> number of distinct member TC numbers supporting it)
    fam_go: dict[str, collections.Counter] = collections.defaultdict(collections.Counter)
    fam_members: dict[tuple[str, str], set[str]] = collections.defaultdict(set)
    fam_name: dict[str, str] = {}
    for go_id, tc, name in rows:
        if go_id not in mf:
            continue
        fam = ".".join(tc.split(".")[:3])
        fam_name.setdefault(fam, name)
        fam_members[(fam, go_id)].add(tc)
    for (fam, go_id), members in fam_members.items():
        fam_go[fam][go_id] = len(members)

    all_go = sorted({g for c in fam_go.values() for g in c})
    labels = go_labels(all_go)

    mappings: list[dict] = []
    for fam in sorted(fam_go):
        # in_scope = every non-obsolete transporter-activity MF term TCDB associates with the
        # family, BEFORE the support threshold. Mono-specificity must be judged on this set:
        # judging it on the filtered set would call a family "mono-specific" (and emit the
        # propagate-safe exactMatch) merely because its other substrate terms each happen to have
        # a single supporting member -- manufacturing the over-generality this project warns about.
        in_scope = {
            g: n for g, n in fam_go[fam].items()
            if g in labels and not labels[g]["obsolete"]
        }
        surviving = {g: n for g, n in in_scope.items() if n >= min_support}
        if not surviving:
            continue
        mono = len(in_scope) == 1
        below = len(in_scope) - len(surviving)
        for go_id, support in sorted(surviving.items(), key=lambda kv: (-kv[1], kv[0])):
            pred = ("skos:exactMatch", "exact match") if mono else ("skos:narrowMatch", "narrow match")
            kind = "mono-specific" if mono else "poly-specific (term applies to a subfamily)"
            if below:
                kind += (f"; {below} further MF term(s) for this family fell below the "
                         f"--min-support {min_support} threshold")
            mappings.append({
                "subject_id": f"TC:{fam}",
                "subject_label": fam_name.get(fam, "").strip(),
                "predicate_id": pred[0],
                "predicate_label": pred[1],
                "object_id": go_id,
                "object_label": labels[go_id]["name"],
                "mapping_justification": "semapv:CompositeMatching",
                "comment": (
                    f"Derived: TCDB go.py associates TC family {fam} members with {go_id} "
                    f"({support} distinct 5-level member(s)). {kind}. "
                    "Not in any external2go mapping."
                ),
            })
    return mappings


HEADER = """# TCDB family -> GO molecular-function mapping (SSSOM) -- GENERATED, DO NOT HAND-EDIT
#
# Auto-generated by projects/TCDB/build_tc2go.py. Each row is the live join of TCDB's own per-protein
# GO dump (https://tcdb.org/cgi-bin/projectv/public/go.py) with the QuickGO ontology, keeping ONLY GO
# terms in the transporter-activity molecular-function closure (is_a GO:0022857 / GO:0005215) and
# aggregating 5-level TC numbers up to the TC family (3-level). No row is hand-typed; the hand-curated,
# review-backed rows live in the sibling tc2go.sssom.yaml.
#
# exactMatch = family maps to a single transporter-activity MF term (mono-specific: one GO term fits the family);
# narrowMatch = family maps to several (poly-specific; the term applies to a subfamily). These are
# TCDB's assertions filtered for aspect + obsolescence, NOT a curated pipeline -- review before use.
"""


def to_sssom(mappings: list[dict]) -> dict:
    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/",
            "sssom": "https://w3id.org/sssom/",
        },
        "mapping_set_id": "https://w3id.org/ai4curation/ai-gene-review/mappings/tc2go-generated",
        "mapping_set_title": "TCDB family -> GO molecular function (generated from TCDB go.py x QuickGO)",
        "mapping_set_description": (
            "Auto-generated TC-family -> GO molecular-function mappings, distilled from TCDB's own "
            "per-protein GO dump (go.py) by keeping only transporter-activity MF terms (is_a "
            "GO:0022857/GO:0005215) and aggregating to the TC family. exactMatch = mono-specific "
            "family; narrowMatch = poly-specific family (term applies to a subfamily). Companion to "
            "the hand-curated tc2go.sssom.yaml seed."
        ),
        "license": "https://creativecommons.org/licenses/by/4.0/",
        "creator_label": ["AI Gene Review project (build_tc2go.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("-o", "--output", type=Path, default=HERE / "tc2go.generated.sssom.yaml")
    ap.add_argument("--min-support", type=int, default=2, help="min distinct member TC numbers per (family,GO)")
    args = ap.parse_args(argv)
    mappings = build(args.min_support)
    body = yaml.safe_dump(to_sssom(mappings), sort_keys=False, default_flow_style=False, width=100, allow_unicode=True)
    args.output.write_text(HEADER + body)
    fams = len({m["subject_id"] for m in mappings})
    exact = sum(1 for m in mappings if m["predicate_id"] == "skos:exactMatch")
    print(f"# wrote {len(mappings)} rows ({fams} families; {exact} exactMatch, {len(mappings)-exact} narrowMatch) -> {args.output}")
    return 0


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