"""Generate per-species codon-usage tables from real Kazusa CUTG data.

Reads ``core/data/cutg_plant_subset.spsum`` (the nuclear codon counts for the 18
supported crops, taken verbatim from the Kazusa Codon Usage Tabulated from
GenBank plant division) and writes ``core/codon_usage_tables.py`` containing:

    CODON_USAGE_TABLES   {species: {aa: {codon: fraction}}}   real per-AA usage
    SPECIES_CDS_COUNT    {species: int}                       CDS compiled (provenance)

Every species therefore uses its OWN measured codon usage — no cross-species
proxies. Run from the ``cool/`` directory:

    python -m tools.build_codon_usage

This keeps the data reproducible: the spsum subset is committed, and re-running
the generator reproduces the table module byte-for-byte.
"""
import os

from core.codons import CODON_TO_AMINO_ACID

HERE = os.path.dirname(os.path.abspath(__file__))
CORE = os.path.normpath(os.path.join(HERE, "..", "core"))
SUBSET = os.path.join(CORE, "data", "cutg_plant_subset.spsum")
OUT = os.path.join(CORE, "codon_usage_tables.py")

# CUTG / SPSUM codon order (RNA), per the SPSUM_LABEL file shipped with CUTG.
ORDER_RNA = ("CGA CGC CGG CGU AGA AGG CUA CUC CUG CUU UUA UUG UCA UCC UCG UCU AGC AGU "
             "ACA ACC ACG ACU CCA CCC CCG CCU GCA GCC GCG GCU GGA GGC GGG GGU GUA GUC GUG GUU "
             "AAA AAG AAC AAU CAA CAG CAC CAU GAA GAG GAC GAU UAC UAU UGC UGU UUC UUU "
             "AUA AUC AUU AUG UGG UAA UAG UGA").split()
ORDER = [c.replace("U", "T") for c in ORDER_RNA]
assert len(ORDER) == 64

# amino-acid order for stable, readable output
AA_ORDER = list("ACDEFGHIKLMNPQRSTVWY") + ["_"]


def parse_subset(path):
    species, counts, names = {}, {}, {}
    cur = None
    with open(path, errors="replace") as fh:
        lines = [ln.rstrip("\n") for ln in fh]
    i = 0
    while i < len(lines):
        ln = lines[i]
        if ln.startswith("# species="):
            cur = ln.split("=", 1)[1].strip()
        elif cur and ln and ln[0].isdigit() and ":" in ln:
            names[cur] = ln.split(":", 1)[1].rsplit(":", 1)[0].strip()
            counts[cur] = int(ln.rsplit(":", 1)[1])
            nums = list(map(int, lines[i + 1].split()))
            assert len(nums) == 64, f"{cur}: expected 64 codon counts"
            species[cur] = nums
            i += 1
            cur = None
        i += 1
    return species, counts, names


def to_usage(raw):
    by_aa = {}
    for codon, n in zip(ORDER, raw):
        by_aa.setdefault(CODON_TO_AMINO_ACID[codon], {})[codon] = n
    out = {}
    for aa, cod in by_aa.items():
        tot = sum(cod.values())
        out[aa] = {c: (round(v / tot, 3) if tot else 0.0) for c, v in cod.items()}
    return out


def fmt_table(usage):
    rows = []
    for aa in AA_ORDER:
        if aa not in usage:
            continue
        inner = ", ".join(f"'{c}': {f:.3f}" for c, f in usage[aa].items())
        rows.append(f"        '{aa}': {{{inner}}},")
    return "\n".join(rows)


def main():
    species, counts, names = parse_subset(SUBSET)
    order = ['arabidopsis', 'rice', 'maize', 'tomato', 'soybean', 'wheat', 'barley',
             'sorghum', 'potato', 'cassava', 'tobacco', 'grape', 'cotton', 'sugarcane',
             'canola', 'banana', 'peanut', 'sunflower']
    lines = ['"""Per-species codon-usage tables — GENERATED, do not edit by hand.',
             "",
             "Source: Kazusa CUTG (Codon Usage Tabulated from GenBank), plant division.",
             "Each table is the measured nuclear codon usage for that species' own taxid —",
             "no cross-species proxies. Regenerate with:  python -m tools.build_codon_usage",
             "",
             "SPECIES_CDS_COUNT records how many CDS were compiled per species (provenance /",
             "confidence): larger = more reliable codon statistics.",
             '"""',
             "",
             "CODON_USAGE_TABLES = {"]
    for sp in order:
        u = to_usage(species[sp])
        lines.append(f"    # {names[sp]} (taxid in subset; {counts[sp]:,} CDS compiled)")
        lines.append(f"    '{sp}': {{")
        lines.append(fmt_table(u))
        lines.append("    },")
    lines.append("}")
    lines.append("")
    lines.append("SPECIES_CDS_COUNT = {")
    for sp in order:
        lines.append(f"    '{sp}': {counts[sp]},")
    lines.append("}")
    lines.append("")
    with open(OUT, "w") as fh:
        fh.write("\n".join(lines))
    print("wrote", os.path.relpath(OUT), "—", len(order), "species")
    for sp in order:
        print(f"  {sp:12} {counts[sp]:>7,} CDS")


if __name__ == "__main__":
    main()