#!/usr/bin/env python3
"""
make_bundle.py — the ONE-TIME build step behind this bundle. Provenance only.

It never runs on visualizationhub.org (Gosling Designer). It ran once, here, to produce the two
things that plain Gosling cannot derive at render time:

  1. data_derived/eqtl_ld_variants.csv
     Gosling has no join. The eQTL-LD view needs per-variant geometry (O1e_credible_variants.tsv:
     var_pos / var_pip / in_ccre) AND the real 1000G r^2 to the lead (O1f_ld_pairs.tsv: r2_cohort /
     r2_gold_eur), which live in two files keyed on (link_key, varID == member_var). This joins them
     and precomputes the `cat` colour bin that the spec's nominal colour channel reads.

  2. build/program_gene_lists.json
     Program membership lives in a PARQUET (O1_program_gene_map.parquet), which Gosling cannot read.
     Rather than ship a data file, each program's gene list is inlined into the program specs as a
     {"type":"filter","field":"gene","oneOf":[...]} transform. This emits those lists.

Everything else in data/ is an untouched copy of an O-script output.

Deps: pyarrow (parquet read) + stdlib csv. No pandas.
"""
from __future__ import annotations

import csv
import json
import sys
from pathlib import Path

SRC = Path("/net/users/voy/nmf/data_itcr/cnTFeffects/reanalysis0612tecounts")
TBL = SRC / "results/O1_ccre_gene_map/tables"
O1F = SRC / "genotype_ld/results/O1f_eqtl_ld_1kg/tables"
OUT = Path(__file__).resolve().parent.parent

# The r^2 colour bins — identical thresholds/labels to z_eqtl_ld_fork.py's `_cat()`, so the platform
# view reads the same as the marimo fork. Order matters: it is the spec's colour `domain` order.
CATS = ["in our cCRE", "out-of-cCRE · r²≥0.8", "out-of-cCRE · 0.2–0.8",
        "out-of-cCRE · r²<0.2", "out-of-cCRE · not in 1000G"]


def _cat(in_ccre: bool, r2) -> str:
    if in_ccre:
        return CATS[0]
    if r2 is None:
        return CATS[4]
    if r2 >= 0.8:
        return CATS[1]
    if r2 >= 0.2:
        return CATS[2]
    return CATS[3]


def _truthy(v: str) -> bool:
    return str(v).strip().lower() in {"true", "1", "t", "yes"}


def _num(v: str):
    try:
        return float(v)
    except (TypeError, ValueError):
        return None


def build_eqtl_ld_variants() -> int:
    """Join O1e_credible_variants ⋈ O1f_ld_pairs on (link_key, varID==member_var); add `cat`."""
    # ---- index the r^2 pairs: (link_key, member_var) -> r2 columns
    r2_by_key: dict[tuple[str, str], dict] = {}
    with (O1F / "O1f_ld_pairs.tsv").open(newline="") as fh:
        for row in csv.DictReader(fh, delimiter="\t"):
            r2_by_key[(row["link_key"], row["member_var"])] = row

    out_path = OUT / "data_derived" / "eqtl_ld_variants.csv"
    out_path.parent.mkdir(parents=True, exist_ok=True)

    # Emitted as TAB-separated (the spec declares "separator": "\t"), matching every other file here.
    cols = ["link_key", "gene", "class", "ccre_chrom", "ccre_start", "ccre_end",
            "varID", "var_chrom", "var_pos", "var_pip", "in_ccre",
            "r2_cohort", "r2_gold_eur", "cat"]

    n = 0
    with (TBL / "O1e_credible_variants.tsv").open(newline="") as fin, \
            out_path.open("w", newline="") as fout:
        w = csv.DictWriter(fout, fieldnames=cols, delimiter="\t", extrasaction="ignore")
        w.writeheader()
        for row in csv.DictReader(fin, delimiter="\t"):
            pair = r2_by_key.get((row["link_key"], row["varID"]))
            r2c = _num(pair["r2_cohort"]) if pair else None
            r2g = _num(pair["r2_gold_eur"]) if pair else None
            in_ccre = _truthy(row["in_ccre"])
            out = {c: row.get(c, "") for c in cols}
            out["in_ccre"] = "true" if in_ccre else "false"
            out["r2_cohort"] = "" if r2c is None else f"{r2c:.4f}"
            out["r2_gold_eur"] = "" if r2g is None else f"{r2g:.4f}"
            out["cat"] = _cat(in_ccre, r2c)
            w.writerow(out)
            n += 1
    return n


def build_program_gene_lists(min_arcs: int = 1) -> dict:
    """program -> sorted gene list, restricted to genes that actually carry arcs in O1g_arcs.tsv."""
    import pyarrow.parquet as pq

    with (SRC / "results/O1_ccre_gene_map/browser/O1g_arcs.tsv").open(newline="") as fh:
        arc_genes = {r["gene"] for r in csv.DictReader(fh, delimiter="\t")}

    tab = pq.read_table(TBL / "O1_program_gene_map.parquet", columns=["program", "gene", "member"])
    progs: dict[str, set] = {}
    for program, gene, member in zip(tab.column("program").to_pylist(),
                                     tab.column("gene").to_pylist(),
                                     tab.column("member").to_pylist()):
        if member and gene in arc_genes:
            progs.setdefault(str(program), set()).add(str(gene))

    lists = {p: sorted(g) for p, g in progs.items() if len(g) >= min_arcs}
    return dict(sorted(lists.items(), key=lambda kv: (-len(kv[1]), kv[0])))


def main() -> int:
    n = build_eqtl_ld_variants()
    print(f"data_derived/eqtl_ld_variants.csv        {n:,} variant rows")

    lists = build_program_gene_lists()
    p = OUT / "build" / "program_gene_lists.json"
    p.write_text(json.dumps(lists, indent=1))
    print(f"build/program_gene_lists.json            {len(lists)} programs, "
          f"{sum(len(v) for v in lists.values()):,} program-gene memberships")
    for name, genes in list(lists.items())[:5]:
        print(f"    {name:<28} {len(genes):>4} genes")
    return 0


if __name__ == "__main__":
    sys.exit(main())