"""Join ComplexPortal-assigned GOA rows to the action given in each gene review.

Usage (from repo root):
  uv run python projects/COMPLEXPORTAL/scripts/cp_review_join.py OUT.json [REJECTED.yaml]

Each CP row (always positive: CP does not assert NOT) is matched to review annotations on
(GO term, evidence code, reference), skipping review annotations marked `negated: true` so a CP
row never inherits the verdict on a NOT row with the same key. Where a review records a
`qualifier`, it must agree with the GOA qualifier. If more than one candidate remains, the first
is used and the row is counted as ambiguous.

OUT.json holds every CP row with its review action. If REJECTED.yaml is given, rows whose action is
MARK_AS_OVER_ANNOTATED, MODIFY or REMOVE are also written there with the review summary.
"""
import collections
import csv
import glob
import json
import os
import sys

import yaml

BAD = ("REMOVE", "MODIFY", "MARK_AS_OVER_ANNOTATED")


def join():
    rows, no_review, ambiguous = [], 0, 0
    for tsv in sorted(glob.glob("genes/*/*/*-goa.tsv")):
        d = os.path.dirname(tsv)
        gene, org = os.path.basename(d), d.split("/")[1]
        cp = [r for r in csv.DictReader(open(tsv), delimiter="\t") if r.get("ASSIGNED BY") == "ComplexPortal"]
        if not cp:
            continue
        rev = os.path.join(d, f"{gene}-ai-review.yaml")
        if not os.path.exists(rev):
            no_review += len(cp)
            continue
        y = yaml.safe_load(open(rev)) or {}
        idx = collections.defaultdict(list)
        for a in y.get("existing_annotations") or []:
            if a.get("negated"):
                continue
            key = ((a.get("term") or {}).get("id"), a.get("evidence_type"), a.get("original_reference_id"))
            idx[key].append(a)
        for r in cp:
            hits = idx.get((r["GO TERM"], r["GO EVIDENCE CODE"], r["REFERENCE"]), [])
            hits = [a for a in hits if not a.get("qualifier") or a["qualifier"] == r["QUALIFIER"]]
            if len(hits) > 1:
                ambiguous += 1
            review = (hits[0].get("review") or {}) if hits else {}
            rows.append(dict(
                org=org, gene=gene, qualifier=r["QUALIFIER"], term=r["GO TERM"], name=r["GO NAME"],
                aspect=r["GO ASPECT"], ev=r["GO EVIDENCE CODE"], ref=r["REFERENCE"],
                action=review.get("action") or "UNMATCHED", summary=review.get("summary") or "",
            ))
    return rows, no_review, ambiguous


def write_rejected(rows, path):
    order = {a: i for i, a in enumerate(BAD)}
    bad = sorted((r for r in rows if r["action"] in order),
                 key=lambda r: (order[r["action"]], r["org"], r["gene"], r["term"]))
    out = [{
        "organism": r["org"], "gene": r["gene"], "qualifier": r["qualifier"],
        "term": {"id": r["term"], "label": r["name"]}, "evidence": r["ev"], "reference": r["ref"],
        "review_action": r["action"], "review_summary": " ".join(r["summary"].split()),
        "review_file": f"genes/{r['org']}/{r['gene']}/{r['gene']}-ai-review.yaml",
    } for r in bad]
    doc = {
        "description": "ComplexPortal-assigned GOA subunit rows that the existing gene review rejected "
                       "(MARK_AS_OVER_ANNOTATED, MODIFY or REMOVE). Generated by "
                       "scripts/cp_review_join.py; review_summary is copied from the gene review.",
        "count": len(out),
        "rows": out,
    }
    yaml.safe_dump(doc, open(path, "w"), sort_keys=False, width=200, allow_unicode=True)


def main():
    rows, no_review, ambiguous = join()
    json.dump(rows, open(sys.argv[1], "w"))
    if len(sys.argv) > 2:
        write_rejected(rows, sys.argv[2])
    print("CP rows in reviewed genes", len(rows), "| rows in genes without review", no_review,
          "| ambiguous matches", ambiguous, "| genes", len({(r["org"], r["gene"]) for r in rows}))
    print(collections.Counter(r["action"] for r in rows).most_common())
    for key in ("ev", "aspect", "org"):
        for v in sorted({r[key] for r in rows}):
            print(key, v, collections.Counter(r["action"] for r in rows if r[key] == v).most_common())


if __name__ == "__main__":
    main()