#!/usr/bin/env python3
"""Select the eggNOG-mapper evaluation cohort from completed AIGR gene reviews.

A review is eligible when every live ``existing_annotations`` row carries a
non-PENDING action, it has at least ``MIN_ANNOTATIONS`` such rows, and its
cached ``*-uniprot.txt`` holds a sequence and an NCBI taxon id. Eligible reviews
are sampled per organism directory with a fixed seed (``PER_ORGANISM`` from each
of the larger organisms; the remaining organisms are pooled into one stratum of
the same size), so the cohort is reproducible and no single proteome dominates.

Outputs (next to this script):
  cohort.yaml   one record per protein: organism dir, gene, accession, taxon, review path
  cohort.fasta  sequences, headers ``>ACCESSION`` (input to emapper)

Run:  python3 projects/EGGNOG_MAPPER_EVALUATION/build_cohort.py
"""
from __future__ import annotations

import glob
import os
import random
import re

import yaml

LOADER = getattr(yaml, "CSafeLoader", yaml.SafeLoader)
HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.abspath(os.path.join(HERE, "..", ".."))

SEED = 20261014
PER_ORGANISM = 40
MIN_ANNOTATIONS = 3
# Organism directories with enough completed reviews to form their own stratum.
STRATA = ["human", "PSEPK", "yeast", "DROME", "worm", "ARATH", "SCHPO",
          "DANRE", "mouse", "ECOLI", "DICDI", "rat"]


def read_uniprot(path: str):
    acc = taxon = None
    seq = []
    in_seq = False
    with open(path) as fh:
        for line in fh:
            if line.startswith("AC   ") and acc is None:
                acc = line[5:].split(";")[0].strip()
            elif line.startswith("OX   "):
                m = re.search(r"NCBI_TaxID=(\d+)", line)
                if m:
                    taxon = int(m.group(1))
            elif line.startswith("SQ   "):
                in_seq = True
            elif line.startswith("//"):
                break
            elif in_seq:
                seq.append(line.strip().replace(" ", ""))
    return acc, taxon, "".join(seq)


def eligible(review_path: str):
    try:
        with open(review_path) as fh:
            doc = yaml.load(fh, Loader=LOADER)
    except Exception:  # noqa: BLE001 - malformed files are simply not eligible
        return None
    if not isinstance(doc, dict):
        return None
    rows = [a for a in doc.get("existing_annotations") or []
            if isinstance(a, dict) and not a.get("retired")]
    actions = [((a.get("review") or {}).get("action") or "PENDING") for a in rows]
    if len(rows) < MIN_ANNOTATIONS or "PENDING" in actions:
        return None
    return doc


def main() -> None:
    by_org: dict[str, list[dict]] = {}
    for review in sorted(glob.glob(os.path.join(ROOT, "genes", "*", "*", "*-ai-review.yaml"))):
        org, gene = review.split(os.sep)[-3:-1]
        uniprot = os.path.join(os.path.dirname(review), f"{gene}-uniprot.txt")
        if not os.path.exists(uniprot):
            continue
        doc = eligible(review)
        if doc is None:
            continue
        acc, taxon, seq = read_uniprot(uniprot)
        if not (acc and taxon and seq):
            continue
        if str(doc.get("id", "")).split(":")[-1] not in ("", acc):
            continue  # review id disagrees with the cached UniProt entry; skip
        stratum = org if org in STRATA else "other"
        by_org.setdefault(stratum, []).append({
            "organism_dir": org, "gene": gene, "accession": acc, "taxon": taxon,
            "review": os.path.relpath(review, ROOT), "length": len(seq), "_seq": seq})

    rng = random.Random(SEED)
    cohort = []
    for stratum in sorted(by_org):
        pool = by_org[stratum]
        pick = pool if len(pool) <= PER_ORGANISM else rng.sample(pool, PER_ORGANISM)
        for rec in sorted(pick, key=lambda r: (r["organism_dir"], r["gene"])):
            cohort.append(dict(rec, stratum=stratum))
        print(f"{stratum:8s} eligible={len(pool):5d} sampled={len(pick)}")

    # One record per accession (two review folders can share one, e.g. the
    # HOST__ACC and paralog conventions); keep the first, for both outputs.
    seen = set()
    unique = []
    for rec in cohort:
        if rec["accession"] not in seen:
            seen.add(rec["accession"])
            unique.append(rec)
    cohort = unique
    with open(os.path.join(HERE, "cohort.fasta"), "w") as fa:
        for rec in cohort:
            fa.write(f">{rec['accession']}\n")
            s = rec["_seq"]
            for i in range(0, len(s), 60):
                fa.write(s[i:i + 60] + "\n")
    out = [{k: v for k, v in r.items() if not k.startswith("_")} for r in cohort]
    with open(os.path.join(HERE, "cohort.yaml"), "w") as fh:
        fh.write("# Generated by build_cohort.py; do not edit by hand.\n")
        yaml.safe_dump({"seed": SEED, "per_organism": PER_ORGANISM,
                        "min_annotations": MIN_ANNOTATIONS, "proteins": out},
                       fh, sort_keys=False)
    print(f"cohort: {len(out)} proteins ({len(seen)} unique accessions)")


if __name__ == "__main__":
    main()