#!/usr/bin/env python
"""Inventory gene-review quotes that use unvetted predictions or labels from UniProt as evidence.

Runs the global best-practice checks in
``ai_gene_review.validation.uniprot_quote_source`` over every gene review and
writes:

* ``uniprot_quote_findings.tsv``: one row per flagged quote, with the review
  element it supports (for an existing annotation, the GO term, its GOA
  reference and the review action) and the quoted statement's evidence;
* ``UNIPROT_QUOTE_AUDIT.md``: counts by check (including protein existence
  cited for an ACCEPT row or a core function), by the unvetted source behind
  the quoted statement (ARBA, imported, electronic GO), by UniProt field and
  by what the quote supports.

Usage::

    uv run python projects/ARBA_REVIEW/uniprot_quote_audit.py
"""

from __future__ import annotations

import collections
import csv
import glob
import re
from pathlib import Path
from typing import Any

import yaml

from ai_gene_review.validation.uniprot_quote_source import (
    UNVETTED_CHECK,
    FIELD_CHECK,
    EXISTENCE_CHECK,
    _describe,
    quote_source_findings,
)

HERE = Path(__file__).parent
ARBA = "GO_REF:0000117"
_ANNOTATION = re.compile(r"^existing_annotations\[(\d+)\]")


def supported_element(data: dict, path: str) -> dict[str, str]:
    """Describe what a quote at *path* supports."""
    m = _ANNOTATION.match(path)
    if m:
        a = data["existing_annotations"][int(m.group(1))]
        return {
            "supports": "existing annotation",
            "term": a.get("term", {}).get("id", ""),
            "label": a.get("term", {}).get("label", ""),
            "annotation_source": "ARBA" if a.get("original_reference_id") == ARBA else a.get("evidence_type", ""),
            "action": (a.get("review") or {}).get("action", ""),
        }
    section = path.split("[", 1)[0].split(".", 1)[0]
    return {"supports": section, "term": "", "label": "", "annotation_source": "", "action": ""}


def main() -> None:
    rows: list[dict[str, Any]] = []
    for f in sorted(glob.glob("genes/*/*/*-ai-review.yaml")):
        data = yaml.load(open(f), Loader=yaml.CSafeLoader) or {}
        for x in quote_source_findings(data):
            evidence = sorted({e for s in x.statements for e in s.evidence})
            sources = sorted({label for s in x.statements for label in s.unvetted_sources})
            rows.append({
                "review": f, "path": x.path, "check": x.check_type, "uniprot_part": _describe(x.statements),
                "unvetted_sources": " + ".join(sources) if x.check_type == UNVETTED_CHECK else "",
                **supported_element(data, x.path),
                "evidence": "; ".join(evidence), "statement": " ".join(x.statements[0].text.split())[:200],
            })
    with open(HERE / "uniprot_quote_findings.tsv", "w", newline="") as fh:
        w = csv.DictWriter(fh, fieldnames=list(rows[0]), delimiter="\t")
        w.writeheader()
        w.writerows(rows)

    lines = [
        "---",
        'title: "UniProt quotes used as evidence: audit"',
        "autolink_gene_symbols: false",
        "---",
        "",
        "# UniProt quotes used as evidence: audit",
        "",
        "Generated by `uniprot_quote_audit.py`; do not edit by hand. Per-quote rows are in",
        "[uniprot_quote_findings.tsv](uniprot_quote_findings.tsv).",
        "",
    ]
    for check, title in [(UNVETTED_CHECK, "Quotes of statements supported only by unvetted predictions"),
                         (FIELD_CHECK, "Quotes of label fields (names, keywords, GO cross-references)"),
                         (EXISTENCE_CHECK, "Protein existence cited to accept a function")]:
        sel = [r for r in rows if r["check"] == check]
        lines += [f"## {title}", "",
                  f"**{len(sel)}** quotes in **{len({r['review'] for r in sel})}** reviews.", ""]
        if check == UNVETTED_CHECK:
            arba = [r for r in sel if "ARBA" in r["unvetted_sources"]]
            lines += [f"ARBA is a source for **{len(arba)}** of these.", "",
                      "| unvetted source | quotes |", "|---|---:|"]
            for src, n in collections.Counter(r["unvetted_sources"] or "other" for r in sel).most_common(12):
                lines.append(f"| {src} | {n} |")
            lines.append("")
        lines += ["| UniProt part | quotes |", "|---|---:|"]
        for part, n in collections.Counter(r["uniprot_part"] for r in sel).most_common(15):
            lines.append(f"| {part} | {n} |")
        lines += ["", "| what the quote supports | quotes |", "|---|---:|"]
        support = collections.Counter(
            f"existing annotation from {r['annotation_source']}, {r['action'] or 'no action'}"
            if r["supports"] == "existing annotation" else r["supports"] for r in sel)
        for s, n in support.most_common(15):
            lines.append(f"| {s} | {n} |")
        lines.append("")
    (HERE / "UNIPROT_QUOTE_AUDIT.md").write_text("\n".join(lines))
    print("\n".join(lines[6:]))


if __name__ == "__main__":
    main()