#!/usr/bin/env python3
"""
gen_specs.py — emits the six Gosling view-config JSONs in ../specs/. Provenance only; never runs on
the platform. This is a faithful port of the marimo forks' builders (reference/z_gb_specs.py and
reference/z_ev_specs.py) from the Python `gosling` API to plain Gosling grammar JSON.

Every invariant from those modules is preserved:
  · rho -> viridis, domain pinned to [0.28, 0.67], `range` deliberately OMITTED (gosling ignores a
    custom range on a quantitative colour channel and supplies its own perceptually-uniform ramp).
  · links are coloured on `stroke`, NOT `color` (a thin withinLink shows only its stroke; an unset
    stroke defaults to black -- this is why the old circular chords were uniformly black).
  · every track in a view shares one `linkingId`, so pan/zoom is linked.
  · lollipops = a bar STEM (0 -> PIP) overlaid with a point HEAD, on a left PIP axis [0, 1].

What used to be pandas now lives in the spec as dataTransform (filter / replace) or as inlined
constants (the 24-row hg38 ideogram, the per-program gene list), so the files in ../data/ stay
byte-identical copies of the O-script outputs.
"""
from __future__ import annotations

import csv
import json
from pathlib import Path

OUT = Path(__file__).resolve().parent.parent
SPECS = OUT / "specs"
DATA = OUT / "data"

# ---------------------------------------------------------------------------------------------
# The one string you rewire after uploading the files to your Gosling Designer workspace.
# ---------------------------------------------------------------------------------------------
BASE = "REPLACE_WITH_WORKSPACE_URL"

ARCS = f"{BASE}/O1g_arcs.tsv"
CVAR = f"{BASE}/O1e_credible_variants.tsv"
CRISPR = f"{BASE}/O1d_crispr_pairs.tsv"
EQTLLD = f"{BASE}/eqtl_ld_variants.csv"

# ---- palette (identical to the forks' shared dict `C`) --------------------------------------
C = dict(NAVY="#10243F", TEAL="#107A73", GOLD="#B87A0E", GREEN="#1E7D46", RED="#AF3A2E",
         SLATE="#4A4E69", MUTE="#6B7785", GREY="#C9D3DD", BG="#F5F7FA", INK="#1B222B")

RHO_DOMAIN = [0.28, 0.67]
BAND_DARK, BAND_LIGHT = "#3B4A5A", "#8CA0B3"

HG38 = {"chr1": 248956422, "chr2": 242193529, "chr3": 198295559, "chr4": 190214555,
        "chr5": 181538259, "chr6": 170805979, "chr7": 159345973, "chr8": 145138636,
        "chr9": 138394717, "chr10": 133797422, "chr11": 135086622, "chr12": 133275309,
        "chr13": 114364328, "chr14": 107043718, "chr15": 101991189, "chr16": 90338345,
        "chr17": 83257441, "chr18": 80373285, "chr19": 58617616, "chr20": 64444167,
        "chr21": 46709983, "chr22": 50818468, "chrX": 156040895, "chrY": 57227415}
CHR_ORDER = list(HG38)

WIDTH = 860

# eQTL-LD r^2 colour bins — order and colours identical to z_eqtl_ld_fork.py
LD_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"]
LD_RANGE = [C["GOLD"], C["RED"], C["SLATE"], C["GREY"], C["MUTE"]]

# CRISPR confusion categories — from z_ev_specs.CRISPR_CATS
CRISPR_CATS = ["TP", "FP", "FN", "TN"]
CRISPR_RANGE = [C["GREEN"], C["RED"], C["GREY"], C["SLATE"]]


# ===============================================================================================
# small helpers
# ===============================================================================================
def rows(path, delim="\t"):
    with Path(path).open(newline="") as fh:
        return list(csv.DictReader(fh, delimiter=delim))


def arc_data(genomic_fields):
    """The O1g arc table as a Gosling `csv` source. All four coordinate columns are declared so the
    same source drives the arc track (element_mid->tss) and the cCRE rect track (start->end)."""
    return {"type": "csv", "url": ARCS, "separator": "\t",
            "chromosomeField": "ccre_chrom", "genomicFields": genomic_fields}


def rho_color(legend, field="rho"):
    """The one rho scale shared by arcs, cCRE rects and skyline so a single legend reads for all.
    `range` intentionally omitted -- see module docstring."""
    return {"field": field, "type": "quantitative", "domain": RHO_DOMAIN, "legend": legend}


def present_chroms(chroms):
    seen = set(chroms)
    return [c for c in CHR_ORDER if c in seen]


def ideogram_data(present):
    """The 24 hg38 chromosome rects, INLINED (24 rows -- a constant, not a data file)."""
    return {"type": "json",
            "values": [{"chrom": c, "s": 1, "e": HG38[c], "chr_label": c} for c in present],
            "chromosomeField": "chrom", "genomicFields": ["s", "e"]}


def chr_bands(present, legend=False):
    rng = [BAND_DARK if i % 2 == 0 else BAND_LIGHT for i, _ in enumerate(present)]
    return {"field": "chr_label", "type": "nominal", "domain": list(present),
            "range": rng, "legend": legend}


def ideogram_track(present, linking_id, axis="top", height=12, width=WIDTH):
    return {"data": ideogram_data(present), "mark": "rect",
            "x": {"field": "s", "type": "genomic", "axis": axis, "linkingId": linking_id},
            "xe": {"field": "e", "type": "genomic"},
            "color": chr_bands(present), "width": width, "height": height}


def write(name, spec):
    p = SPECS / name
    p.parent.mkdir(parents=True, exist_ok=True)
    p.write_text(json.dumps(spec, indent=2, ensure_ascii=False))
    return p


# ===============================================================================================
# A · LOCUS — one gene: rho arcs, rho cCREs, gene anchor, PIP lollipops
# ===============================================================================================
def build_locus(gene, arcs_rows, cvar_rows):
    mine = [r for r in arcs_rows if r["gene"] == gene]
    chrom = mine[0]["ccre_chrom"]
    tss = int(float(mine[0]["tss"]))
    vs = [r for r in cvar_rows if r["gene"] == gene]

    pts = ([int(float(r["ccre_start"])) for r in mine] + [int(float(r["ccre_end"])) for r in mine]
           + [tss] + [int(float(v["var_pos"])) for v in vs])
    lo, hi = min(pts), max(pts)
    pad = max(1500, int((hi - lo) * 0.08))
    lo, hi = lo - pad, hi + pad

    LID = "locusdetail"
    dom = {"chromosome": chrom, "interval": [lo, hi]}
    gene_filter = [{"type": "filter", "field": "gene", "oneOf": [gene]}]

    x_arc = {"field": "element_mid", "type": "genomic", "domain": dom, "linkingId": LID}

    arcs_t = {
        "data": arc_data(["element_mid", "tss", "ccre_start", "ccre_end"]),
        "dataTransform": gene_filter, "mark": "withinLink",
        "x": x_arc, "xe": {"field": "tss", "type": "genomic"},
        "stroke": rho_color(True), "color": rho_color(False),
        "strokeWidth": {"value": 1.6}, "opacity": {"value": 0.6},
        "width": WIDTH, "height": 110,
    }
    ccre_t = {
        "data": arc_data(["element_mid", "tss", "ccre_start", "ccre_end"]),
        "dataTransform": gene_filter, "mark": "rect",
        "x": {"field": "ccre_start", "type": "genomic", "domain": dom, "linkingId": LID},
        "xe": {"field": "ccre_end", "type": "genomic"},
        "color": rho_color(False), "stroke": {"value": C["INK"]}, "strokeWidth": {"value": 0.4},
        "width": WIDTH, "height": 18,
    }

    # GENE ANCHOR: tick + downward caret + label, overlaid (inline 1-row constant)
    anchor_data = {"type": "json",
                   "values": [{"chrom": chrom, "s": tss - 220, "e": tss + 220, "gene": gene}],
                   "chromosomeField": "chrom", "genomicFields": ["s", "e"]}
    x_anchor = {"field": "s", "type": "genomic", "domain": dom, "linkingId": LID}
    xe_anchor = {"field": "e", "type": "genomic"}
    anchor_t = {
        "alignment": "overlay", "data": anchor_data,
        "x": x_anchor, "xe": xe_anchor, "color": {"value": C["INK"]},
        "tracks": [
            {"mark": "rect"},
            {"mark": "triangleBottom", "size": {"value": 8}},
            {"mark": "text", "text": {"field": "gene", "type": "nominal"}, "size": {"value": 11}},
        ],
        "width": WIDTH, "height": 20,
    }

    tracks = [arcs_t, ccre_t, anchor_t]

    # LOLLIPOPS: bar STEM (0->PIP) + point HEAD on a left PIP axis. `evid` is derived IN THE SPEC
    # from in_ccre via a replace transform (the OG file stores Python-style "True"/"False").
    if vs:
        lolli_data = {"type": "csv", "url": CVAR, "separator": "\t",
                      "chromosomeField": "var_chrom", "genomicFields": ["var_pos"]}
        lolli_tf = [
            {"type": "filter", "field": "gene", "oneOf": [gene]},
            {"type": "replace", "field": "in_ccre", "newField": "evid",
             "replace": [{"from": "True", "to": "in cCRE"}, {"from": "False", "to": "LD proxy"}]},
        ]
        x_v = {"field": "var_pos", "type": "genomic", "domain": dom, "linkingId": LID}
        y_v = {"field": "var_pip", "type": "quantitative", "domain": [0, 1], "axis": "left"}
        tracks.append({
            "alignment": "overlay", "data": lolli_data, "dataTransform": lolli_tf,
            "x": x_v, "y": y_v,
            "tracks": [
                {"mark": "bar", "color": {"value": C["GREY"]}, "size": {"value": 1.6}},
                {"mark": "point",
                 "size": {"field": "var_pip", "type": "quantitative",
                          "domain": [0, 1], "range": [3, 9]},
                 "color": {"field": "evid", "type": "nominal",
                           "domain": ["in cCRE", "LD proxy"],
                           "range": [C["RED"], C["GOLD"]], "legend": True}},
            ],
            "width": WIDTH, "height": 92,
        })

    return {
        "title": f"Locus · {gene}",
        "subtitle": (f"{len(mine)} cCRE→gene arcs · {len(vs)} credible variants · "
                     f"ρ-coloured arcs & cCREs, gene anchor, PIP lollipops · {chrom}:{lo:,}-{hi:,}"),
        "assembly": "hg38", "spacing": 8, "layout": "linear",
        "views": [{"tracks": tracks}],
    }, len(mine), len(vs)


# ===============================================================================================
# B · PROGRAM — one NMF program, whole-genome linear (ideogram + arcs + rho skyline) / circular
# ===============================================================================================
def build_program(program, genes, arcs_rows, layout="linear"):
    mine = [r for r in arcs_rows if r["gene"] in set(genes)]
    present = present_chroms(r["ccre_chrom"] for r in mine)
    gene_filter = [{"type": "filter", "field": "gene", "oneOf": sorted(genes)}]
    data = arc_data(["element_mid", "tss", "ccre_start", "ccre_end"])

    if layout == "circular":
        LID = "progring"
        ideo = {"data": ideogram_data(present), "mark": "rect",
                "x": {"field": "s", "type": "genomic", "linkingId": LID},
                "xe": {"field": "e", "type": "genomic"},
                "color": chr_bands(present), "width": 600, "height": 18}
        chords = {"data": data, "dataTransform": gene_filter, "mark": "withinLink",
                  "x": {"field": "element_mid", "type": "genomic", "linkingId": LID},
                  "xe": {"field": "tss", "type": "genomic"},
                  "stroke": rho_color(True), "strokeWidth": {"value": 1.4},
                  "opacity": {"value": 0.7},
                  "style": {"linkStyle": "elliptical"}, "width": 600, "height": 360}
        return {
            "title": f"Program · {program} · circular",
            "subtitle": (f"{len(genes)} member genes · {len(mine)} cCRE→gene arcs as ρ-coloured "
                         f"chords on a chromosome ring (a short cis link reads as a spoke)"),
            "assembly": "hg38", "spacing": 2, "layout": "circular",
            "views": [{"tracks": [ideo, chords]}],
        }, len(mine)

    LID = "progwg"
    ideo = ideogram_track(present, LID, axis="top", height=12)
    arcs_t = {"data": data, "dataTransform": gene_filter, "mark": "withinLink",
              "x": {"field": "element_mid", "type": "genomic", "axis": "none", "linkingId": LID},
              "xe": {"field": "tss", "type": "genomic"},
              "stroke": rho_color(False), "color": rho_color(False),
              "strokeWidth": {"value": 1.2}, "opacity": {"value": 0.5},
              "width": WIDTH, "height": 120}
    sky = {"data": data, "dataTransform": gene_filter, "mark": "point",
           "x": {"field": "element_mid", "type": "genomic", "axis": "none", "linkingId": LID},
           "y": {"field": "rho", "type": "quantitative", "domain": RHO_DOMAIN, "axis": "left"},
           "color": rho_color(True), "size": {"value": 5.5}, "opacity": {"value": 0.85},
           "stroke": {"value": "white"}, "strokeWidth": {"value": 0.3},
           "width": WIDTH, "height": 180}
    return {
        "title": f"Program · {program} · whole genome",
        "subtitle": (f"{len(genes)} member genes · {len(mine)} cCRE→gene arcs. Fully zoomed out a cis "
                     f"arc is sub-pixel, so the ρ skyline carries the overview; zoom into an arm and "
                     f"the arcs resolve (all tracks share one linkingId)."),
        "assembly": "hg38", "spacing": 4, "layout": "linear",
        "views": [{"tracks": [ideo, arcs_t, sky]}],
    }, len(mine)


# ===============================================================================================
# C · ATLAS — every arc, circular Circos (ideogram ring · cCRE ring · chords by chromosome)
# ===============================================================================================
def build_atlas(arcs_rows, size=600):
    present = present_chroms(r["ccre_chrom"] for r in arcs_rows)
    data = arc_data(["element_mid", "tss", "ccre_start", "ccre_end"])

    ideo = {"data": ideogram_data(present), "mark": "rect",
            "x": {"field": "s", "type": "genomic"}, "xe": {"field": "e", "type": "genomic"},
            "color": chr_bands(present), "width": size, "height": 18}
    ccre = {"data": data, "mark": "point",
            "x": {"field": "element_mid", "type": "genomic"},
            "color": chr_bands(present, legend=False),
            "size": {"value": 2.0}, "opacity": {"value": 0.6},
            "width": size, "height": 20}
    chords = {"data": data, "mark": "withinLink",
              "x": {"field": "element_mid", "type": "genomic"},
              "xe": {"field": "tss", "type": "genomic"},
              "stroke": chr_bands(present), "strokeWidth": {"value": 0.9},
              "opacity": {"value": 0.55},
              "style": {"linkStyle": "elliptical"}, "width": size, "height": int(size * 0.52)}
    return {
        "title": "Atlas · every cCRE→gene arc",
        "subtitle": (f"{len(arcs_rows)} arcs across {len(present)} chromosomes · outer ideogram ring, "
                     f"a ring of every cCRE position, and the chords coloured by chromosome"),
        "assembly": "hg38", "spacing": 2, "layout": "circular",
        "views": [{"tracks": [ideo, ccre, chords]}],
    }, len(arcs_rows)


# ===============================================================================================
# D · eQTL-LD — one link's credible set, variants coloured by REAL 1000G r² to the lead
# ===============================================================================================
def build_eqtl_ld(gene, ld_rows, arcs_rows):
    vs = [r for r in ld_rows if r["gene"] == gene]
    mine = [r for r in arcs_rows if r["gene"] == gene]
    chrom = vs[0]["ccre_chrom"]

    pts = ([int(float(v["var_pos"])) for v in vs]
           + [int(float(v["ccre_start"])) for v in vs] + [int(float(v["ccre_end"])) for v in vs])
    lo, hi = min(pts), max(pts)
    pad = max(1500, int((hi - lo) * 0.12))
    lo, hi = lo - pad, hi + pad
    dom = {"chromosome": chrom, "interval": [lo, hi]}
    LID = "eqtlld"
    gf = [{"type": "filter", "field": "gene", "oneOf": [gene]}]

    ld_data = {"type": "csv", "url": EQTLLD, "separator": "\t",
               "chromosomeField": "ccre_chrom", "genomicFields": ["var_pos", "ccre_start", "ccre_end"]}

    ccre_t = {"data": ld_data, "dataTransform": gf, "mark": "rect",
              "x": {"field": "ccre_start", "type": "genomic", "domain": dom, "linkingId": LID},
              "xe": {"field": "ccre_end", "type": "genomic"},
              "color": {"value": C["TEAL"]}, "width": WIDTH, "height": 26}

    # only advertise categories actually present, so the legend stays honest (as the fork does)
    present_cats = [c for c in LD_CATS if any(v["cat"] == c for v in vs)]
    present_rng = [LD_RANGE[LD_CATS.index(c)] for c in present_cats]

    x_v = {"field": "var_pos", "type": "genomic", "domain": dom, "linkingId": LID}
    y_v = {"field": "var_pip", "type": "quantitative", "domain": [0, 1], "axis": "left"}
    lolli = {
        "alignment": "overlay", "data": ld_data, "dataTransform": gf, "x": x_v, "y": y_v,
        "tracks": [
            {"mark": "bar", "color": {"value": C["GREY"]}, "size": {"value": 1.6}},
            {"mark": "point",
             "size": {"field": "var_pip", "type": "quantitative", "domain": [0, 1], "range": [3, 10]},
             "color": {"field": "cat", "type": "nominal", "domain": present_cats,
                       "range": present_rng, "legend": True}},
        ],
        "width": WIDTH, "height": 120,
    }
    n_in = sum(1 for v in vs if v["cat"] == LD_CATS[0])
    n_hi = sum(1 for v in vs if v["cat"] == LD_CATS[1])
    return {
        "title": f"eQTL LD · {gene} credible set",
        "subtitle": (f"{len(vs)} credible variants ({n_in} inside our cCRE, {n_hi} out-of-cCRE with "
                     f"r²≥0.8) · teal = our cCRE; lollipop height = PIP, colour = real 1000G r² to the "
                     f"lead (cohort model)"),
        "assembly": "hg38", "spacing": 8, "layout": "linear",
        "views": [{"tracks": [ccre_t, lolli]}],
    }, len(vs), len(mine)


# ===============================================================================================
# E · CRISPR — EPCrisprBenchmark pairs by confusion category + CRISPRi effect-size skyline
# ===============================================================================================
def build_crispr(crispr_rows):
    present = present_chroms(r["elem_chrom"] for r in crispr_rows)
    es = [float(r["EffectSize"]) for r in crispr_rows if r["EffectSize"] not in ("", "NA")]
    y_domain = [round(min(es) - 0.02, 2), round(max(es) + 0.02, 2)]

    LID = "crisprwg"
    data = {"type": "csv", "url": CRISPR, "separator": "\t",
            "chromosomeField": "elem_chrom", "genomicFields": ["elem_start", "elem_end", "startTSS"]}
    cats = [c for c in CRISPR_CATS if any(r["category"] == c for r in crispr_rows)]
    rng = [CRISPR_RANGE[CRISPR_CATS.index(c)] for c in cats]

    def cat_color(legend):
        return {"field": "category", "type": "nominal", "domain": cats, "range": rng,
                "legend": legend}

    ideo = ideogram_track(present, LID, axis="top", height=12)
    arcs_t = {"data": data, "mark": "withinLink",
              "x": {"field": "elem_start", "type": "genomic", "axis": "none", "linkingId": LID},
              "xe": {"field": "startTSS", "type": "genomic"},
              "stroke": cat_color(False), "strokeWidth": {"value": 1.1},
              "opacity": {"value": 0.55}, "width": WIDTH, "height": 120}
    sky = {"data": data, "mark": "point",
           "x": {"field": "elem_start", "type": "genomic", "axis": "none", "linkingId": LID},
           "y": {"field": "EffectSize", "type": "quantitative", "domain": y_domain, "axis": "left"},
           "color": cat_color(True), "size": {"value": 5.0}, "opacity": {"value": 0.8},
           "stroke": {"value": "white"}, "strokeWidth": {"value": 0.3},
           "width": WIDTH, "height": 190}
    n = {c: sum(1 for r in crispr_rows if r["category"] == c) for c in cats}
    return {
        "title": "CRISPR · EPCrisprBenchmark element→gene pairs",
        "subtitle": (f"{len(crispr_rows):,} perturbation-tested pairs (K562) · "
                     + " · ".join(f"{c} {n[c]:,}" for c in cats)
                     + " · arcs + CRISPRi effect-size skyline, coloured by confusion category"),
        "assembly": "hg38", "spacing": 4, "layout": "linear",
        "views": [{"tracks": [ideo, arcs_t, sky]}],
    }, len(crispr_rows)


# ===============================================================================================
def main():
    arcs_rows = rows(DATA / "O1g_arcs.tsv")
    cvar_rows = rows(DATA / "O1e_credible_variants.tsv")
    crispr_rows = rows(DATA / "O1d_crispr_pairs.tsv")
    ld_rows = rows(OUT / "data_derived" / "eqtl_ld_variants.csv")
    programs = json.loads((OUT / "build" / "program_gene_lists.json").read_text())

    # ---- LOCUS: the gene with the most cCREs (the richest single-gene picture) ----------------
    by_gene = {}
    for r in arcs_rows:
        by_gene.setdefault(r["gene"], []).append(r)
    cvar_genes = {r["gene"] for r in cvar_rows}
    locus_gene = max((g for g in by_gene if g in cvar_genes),
                     key=lambda g: (len(by_gene[g]), g))
    spec, n_arcs, n_vars = build_locus(locus_gene, arcs_rows, cvar_rows)
    write("spec_locus.json", spec)
    print(f"spec_locus.json               gene={locus_gene}  {n_arcs} arcs  {n_vars} variants")

    # ---- PROGRAM: the program with the most arcs ---------------------------------------------
    prog = max(programs, key=lambda p: len(programs[p]))
    spec, n = build_program(prog, programs[prog], arcs_rows, layout="linear")
    write("spec_program_linear.json", spec)
    print(f"spec_program_linear.json      program={prog}  {len(programs[prog])} genes  {n} arcs")
    spec, n = build_program(prog, programs[prog], arcs_rows, layout="circular")
    write("spec_program_circular.json", spec)
    print(f"spec_program_circular.json    program={prog}  {len(programs[prog])} genes  {n} arcs")

    # ---- ATLAS -------------------------------------------------------------------------------
    spec, n = build_atlas(arcs_rows)
    write("spec_atlas.json", spec)
    print(f"spec_atlas.json               {n} arcs")

    # ---- eQTL-LD: a gene that actually SHOWS the bystander case (an out-of-cCRE r²>=0.8) ------
    ld_by_gene = {}
    for r in ld_rows:
        ld_by_gene.setdefault(r["gene"], []).append(r)

    def score(g):
        v = ld_by_gene[g]
        has_in = any(x["cat"] == LD_CATS[0] for x in v)
        n_hi = sum(1 for x in v if x["cat"] == LD_CATS[1])
        return (has_in and n_hi > 0, n_hi, -abs(len(v) - 40))

    ld_gene = max(ld_by_gene, key=score)
    spec, n_v, n_a = build_eqtl_ld(ld_gene, ld_rows, arcs_rows)
    write("spec_eqtl_ld.json", spec)
    print(f"spec_eqtl_ld.json             gene={ld_gene}  {n_v} credible variants")

    # ---- CRISPR ------------------------------------------------------------------------------
    spec, n = build_crispr(crispr_rows)
    write("spec_crispr.json", spec)
    print(f"spec_crispr.json              {n:,} pairs")


if __name__ == "__main__":
    main()