#!/usr/bin/env python3
"""Characterize HPA cilium-atlas and sperm calls against GOA and our human reviews.

The HPA primary-cilium atlas (Hansen et al. 2025, PMID:41005307) and the sperm
flagellum locations are part of the v25 subcellular file. Most of these calls
are graded Approved or Uncertain and are therefore NOT exported to GOA. This
script reports, for every HPA cilia/sperm call:

  - the HPA grade and mapped GO term
  - whether the gene has a review, and whether GOA already carries the HPA row
  - whether the review already has independent (non-HPA) ciliary/sperm support:
    an accepted/kept/new/modified existing annotation whose label matches
    cilia vocabulary, or cilia vocabulary in core_functions/description.
    This is a crude text match intended for triage, not a gold standard.

Usage (from repo root):
    uv run python projects/HUMAN_PROTEIN_ATLAS/scripts/hpa_cilia.py

Outputs (projects/HUMAN_PROTEIN_ATLAS/data/):
    hpa_cilia_calls.tsv    one row per gene x cilia/sperm location
    hpa_cilia_summary.md   grade x support tables
"""

import csv
import re
from collections import Counter, defaultdict
from pathlib import Path

import yaml

PROJECT = Path("projects/HUMAN_PROTEIN_ATLAS")
HPA_FILE = PROJECT / "data/subcellular_location.tsv"
GRADES = ["Enhanced", "Supported", "Approved", "Uncertain"]
HPA_REFS = {"GO_REF:0000052", "PMID:18029348"}
GROUPS = {
    "primary_cilium": ["Primary cilium", "Primary cilium tip",
                       "Primary cilium transition zone", "Basal body"],
    "sperm": ["Mid piece", "Principal piece", "End piece", "Connecting piece",
              "Flagellar centriole", "Acrosome"],
}
LOC2GROUP = {loc: grp for grp, locs in GROUPS.items() for loc in locs}
CILIA_RE = re.compile(
    r"cili|axonem|basal body|centriol|flagell|sperm|acrosom|intraflagellar|BBSome", re.I)
GO_RE = re.compile(r"^(.*) \((GO:\d{7})\)$")
SUPPORT_ACTIONS = {"ACCEPT", "KEEP_AS_NON_CORE", "NEW", "MODIFY"}


def hpa_goa_terms(goa_file: Path) -> set:
    terms = set()
    if goa_file.exists():
        for row in csv.reader(goa_file.open(), delimiter="\t"):
            if len(row) > 13 and row[13] == "HPA":
                terms.add(row[4])
    return terms


def independent_support(review_file: Path) -> str:
    doc = yaml.safe_load(review_file.read_text()) or {}
    for ann in doc.get("existing_annotations") or []:
        if ann.get("original_reference_id") in HPA_REFS:
            continue
        label = (ann.get("term") or {}).get("label", "")
        action = (ann.get("review") or {}).get("action", "")
        if CILIA_RE.search(label) and action in SUPPORT_ACTIONS:
            return f"annotation:{label}"
    text = yaml.safe_dump(doc.get("core_functions") or []) + (doc.get("description") or "")
    m = CILIA_RE.search(text)
    return f"text:{m.group(0)}" if m else ""


def main():
    rows = []
    for r in csv.DictReader(HPA_FILE.open(), delimiter="\t"):
        grade = {}
        for g in GRADES:
            for loc in filter(None, (r.get(g) or "").split(";")):
                grade[loc] = g
        go = {}
        for item in filter(None, (r.get("GO id") or "").split(";")):
            m = GO_RE.match(item.strip())
            if m:
                go[m.group(1)] = m.group(2)
        gene = r["Gene name"]
        gdir = Path("genes/human") / gene
        review = gdir / f"{gene}-ai-review.yaml"
        reviewed = review.exists()
        goa_terms = hpa_goa_terms(gdir / f"{gene}-goa.tsv") if reviewed else set()
        support = None
        for loc, g in grade.items():
            if loc not in LOC2GROUP:
                continue
            if reviewed and support is None:
                support = independent_support(review)
            gid = go.get(loc, "")
            rows.append({
                "gene": gene, "ensembl": r["Gene"], "group": LOC2GROUP[loc],
                "hpa_location": loc, "go_id": gid, "grade": g,
                "main": loc in r["Main location"].split(";"),
                "reviewed": reviewed,
                "in_goa_from_hpa": gid in goa_terms if reviewed else "",
                "independent_support": support or "" if reviewed else "",
            })

    out = PROJECT / "data"
    with (out / "hpa_cilia_calls.tsv").open("w", newline="") as fh:
        w = csv.DictWriter(fh, fieldnames=list(rows[0]), delimiter="\t", lineterminator="\n")
        w.writeheader()
        w.writerows(rows)

    lines = ["# HPA cilia/sperm calls (generated by scripts/hpa_cilia.py)\n"]
    for grp in GROUPS:
        sub = [r for r in rows if r["group"] == grp]
        genes = {r["gene"] for r in sub}
        lines.append(f"## {grp}: {len(sub)} calls, {len(genes)} genes\n")
        lines.append("| location | " + " | ".join(GRADES) + " |")
        lines.append("|---|" + "---|" * len(GRADES))
        c = Counter((r["hpa_location"], r["grade"]) for r in sub)
        for loc in GROUPS[grp]:
            lines.append(f"| {loc} | " + " | ".join(str(c[(loc, g)]) for g in GRADES) + " |")
        rev = [r for r in sub if r["reviewed"]]
        lines.append(f"\nCalls on reviewed genes: {len(rev)}; already in GOA from HPA: "
                     f"{sum(1 for r in rev if r['in_goa_from_hpa'])}\n")
        # per gene, best grade
        best = {}
        sup = {}
        for r in rev:
            g = r["gene"]
            if g not in best or GRADES.index(r["grade"]) < GRADES.index(best[g]):
                best[g] = r["grade"]
            sup[g] = bool(r["independent_support"])
        lines.append("| best HPA grade | reviewed genes | with independent cilia/sperm support |")
        lines.append("|---|---|---|")
        tally = defaultdict(lambda: [0, 0])
        for g, gr in best.items():
            tally[gr][0] += 1
            tally[gr][1] += sup[g]
        for gr in GRADES:
            n, s = tally[gr]
            if n:
                lines.append(f"| {gr} | {n} | {s} ({100 * s / n:.0f}%) |")
        lines.append("")
    summary = "\n".join(lines) + "\n"
    (out / "hpa_cilia_summary.md").write_text(summary)
    print(summary)


if __name__ == "__main__":
    main()