#!/usr/bin/env python3
"""Build the compact, vendored inputs from the sha256-pinned public sources (run once, outside Docker).
usage: prep.py RAW_DIR BUNDLE_DIR
Writes data/SOURCES.json, data/energy_genes.csv, data/subclass_expression.csv, data/subclasses.csv,
data/cells_<brain>.npz, data/qc_counts.json. Computes NO cell-to-vessel distance."""
import sys, json, hashlib, pathlib, re
import numpy as np, pandas as pd, h5py

RAW = pathlib.Path(sys.argv[1]); OUT = pathlib.Path(sys.argv[2]) / "data"
P = json.load(open(OUT / "parameters.json"))
B = "https://allen-brain-cell-atlas.s3.us-west-2.amazonaws.com"
SRC = {
 "allen/cell_metadata_with_parcellation_annotation.csv": B + "/metadata/MERFISH-C57BL6J-638850-CCF/20231215/views/cell_metadata_with_parcellation_annotation.csv",
 **{f"zhuang{i}/cell_metadata_with_cluster_annotation.csv": B + f"/metadata/Zhuang-ABCA-{i}/20241115/views/cell_metadata_with_cluster_annotation.csv" for i in range(1, 5)},
 **{f"zhuang{i}/ccf_coordinates.csv": B + f"/metadata/Zhuang-ABCA-{i}-CCF/20230830/ccf_coordinates.csv" for i in range(1, 5)},
 **{f"h5ad/Zhuang-ABCA-{i}-raw.h5ad": B + f"/expression_matrices/Zhuang-ABCA-{i}/20230830/Zhuang-ABCA-{i}-raw.h5ad" for i in range(1, 5)},
 "wmb/precomputed_stats_ABC_revision_230821.h5": B + "/mapmycells/WMB-10X/20240831/precomputed_stats_ABC_revision_230821.h5",
 "wmb/gene.csv": B + "/metadata/WMB-10X/20241115/gene.csv",
 "wmb/cluster.csv": B + "/metadata/WMB-taxonomy/20231215/cluster.csv",
 "wmb/cluster_to_cluster_annotation_membership_pivoted.csv": B + "/metadata/WMB-taxonomy/20231215/views/cluster_to_cluster_annotation_membership_pivoted.csv",
 "ccf/parcellation_to_parcellation_term_membership_acronym.csv": B + "/metadata/Allen-CCF-2020/20230630/views/parcellation_to_parcellation_term_membership_acronym.csv",
 "kegg/mmu00190.txt": "https://rest.kegg.jp/get/mmu00190 (fetched " + (RAW / "kegg/fetched_utc.txt").read_text().strip() + ")",
}

def sha(p):
    h = hashlib.sha256()
    with open(p, "rb") as f:
        for b in iter(lambda: f.read(1 << 24), b""): h.update(b)
    return h.hexdigest()

src = {k: {"url": u, "bytes": (RAW / k).stat().st_size, "sha256": sha(RAW / k)} for k, u in sorted(SRC.items())}
json.dump(src, open(OUT / "SOURCES.json", "w"), indent=1, sort_keys=True)
print("pinned", len(src), "sources", flush=True)

NEURON_CLASS = re.compile(P["neuron_class_regex"])
def is_neuron_class(c): return isinstance(c, str) and bool(NEURON_CLASS.match(c)) and c not in P["excluded_classes"]
def is_imn(s): return isinstance(s, str) and P["immature_subclass_token"] in s.split()

# ---------- energy genes (KEGG mmu00190, locked rule) ----------
txt = (RAW / "kegg/mmu00190.txt").read_text().splitlines()
genes, on = [], False
for ln in txt:
    if ln.startswith("GENE"): on = True
    elif ln[:1].strip(): on = False
    if on:
        tok = ln[4:].split() if ln.startswith("GENE") else ln.split()
        genes.append((tok[1].rstrip(";"), "NCBIGene:" + tok[0]))
keep = sorted({(g, nc) for g, nc in genes if re.match(P["oxphos_symbol_regex"], g) and not re.search(P["oxphos_exclude_regex"], g)})
g = pd.read_csv(RAW / "wmb/gene.csv")
ncbi2id, sym2id = {}, {}
for gid, sy, nc in zip(g.gene_identifier, g.gene_symbol, g.mapped_ncbi_identifier):
    sym2id.setdefault(sy, gid)
    if isinstance(nc, str): ncbi2id.setdefault(nc, gid)
eg = [(sy, nc, "oxphos") for sy, nc in keep] + [(sy, "", "pump") for sy in P["pump_genes"]]
eg = pd.DataFrame(eg, columns=["symbol", "ncbi", "component"])
# KEGG entries are matched to the 10x gene table by NCBI Gene id (10x symbols predate the 2019 Atp5 renaming); pump genes by symbol
eg["gene_identifier"] = [ncbi2id.get(nc) if nc else sym2id.get(sy) for sy, nc in zip(eg.symbol, eg.ncbi)]
id2sym10x = dict(zip(g.gene_identifier, g.gene_symbol))
eg["symbol_10x"] = eg.gene_identifier.map(id2sym10x)
eg = eg[~(eg.gene_identifier.notna() & eg.duplicated("gene_identifier"))].reset_index(drop=True)
eg["in_10x"] = eg.gene_identifier.notna()

f = h5py.File(RAW / "wmb/precomputed_stats_ABC_revision_230821.h5", "r")
cols = json.loads(f["col_names"][()].decode()); c2r = json.loads(f["cluster_to_row"][()].decode())
tt = json.loads(f["taxonomy_tree"][()].decode()); nm = tt["name_mapper"]["CCN20230722_CLUS"]
piv = pd.read_csv(RAW / "wmb/cluster_to_cluster_annotation_membership_pivoted.csv")
cl2sub = dict(zip(piv.cluster, piv.subclass)); cl2cls = dict(zip(piv.cluster, piv["class"]))
clc = pd.read_csv(RAW / "wmb/cluster.csv")
alias2n = dict(zip(clc.cluster_alias, clc.number_of_cells))
name2n = {nm_: alias2n[int(al)] for nm_, al in zip(piv.cluster, piv.cluster_alias)}
col_idx = {c: i for i, c in enumerate(cols)}
eg["in_10x"] = eg.gene_identifier.map(lambda x: x in col_idx if isinstance(x, str) else False)
gi = [col_idx[x] for x in eg.gene_identifier[eg.in_10x]]
S = f["sum"][:, sorted(gi)]; order = {c: k for k, c in enumerate(sorted(gi))}
S = S[:, [order[c] for c in gi]]
rows = []
for lab, r in c2r.items():
    name = nm[lab]["name"]; rows.append((name, cl2sub[name], cl2cls[name], r, name2n[name]))
cl = pd.DataFrame(rows, columns=["cluster", "subclass", "class", "row", "n10x"]).sort_values("row")
# stats.h5 'sum' rows are cluster means (n_cells==1 for every row); weight by 10x cells per cluster (cluster.csv)
assert (f["n_cells"][:] == 1).all()
sub_sum = pd.DataFrame(S[cl.row.values] * cl.n10x.values[:, None], columns=list(eg.symbol[eg.in_10x])).groupby(cl.subclass.values).sum()
sub_n = pd.Series(cl.n10x.values.astype(float)).groupby(cl.subclass.values).sum()
expr = sub_sum.div(sub_n, axis=0)
expr.insert(0, "n_cells_10x", sub_n.astype(int))
expr.index.name = "subclass"
expr.round(8).sort_index().to_csv(OUT / "subclass_expression.csv")
eg.to_csv(OUT / "energy_genes.csv", index=False)
print("energy genes:", int(eg.in_10x.sum()), "of", len(eg), flush=True)

# ---------- divisions ----------
pm = pd.read_csv(RAW / "ccf/parcellation_to_parcellation_term_membership_acronym.csv")
idx2div = dict(zip(pm.parcellation_index, pm.division))

# ---------- cells ----------
def build(df, brain, logcounts=None):
    df = df.copy()
    cls = df["class"].astype(str); sub = df["subclass"].astype(str)
    neur = cls.map(is_neuron_class) & df.division.ne("unassigned") & df.division.notna()
    imn = sub.map(is_imn)
    kind = np.zeros(len(df), np.int8)
    kind[(neur & ~imn & df.hq).values] = 1
    kind[(neur & imn & df.hq).values] = 4
    kind[(sub == P["endothelial_subclass"]).values] = 2
    kind[(sub == P["pericyte_subclass"]).values] = 3
    df["kind"] = kind
    d = df[df.kind > 0].sort_values(["section", "cell_label"], kind="mergesort").reset_index(drop=True)
    secs = sorted(d.section.unique()); s2i = {s: i for i, s in enumerate(secs)}
    out = dict(section_names=np.array(secs, dtype="U"), sec=d.section.map(s2i).astype(np.int16).values,
               x_um=(d.x.values * 1000).astype(np.float32), y_um=(d.y.values * 1000).astype(np.float32),
               kind=d.kind.values.astype(np.int8), subclass=d.subclass.astype(str).to_numpy(dtype="U"), cls=d["class"].astype(str).to_numpy(dtype="U"),
               division=d.division.astype(str).to_numpy(dtype="U"))
    if logcounts is not None:
        out["log10_counts"] = d.cell_label.map(logcounts).astype(np.float32).values
    np.savez_compressed(OUT / f"cells_{brain}.npz", **out)
    q = {"n_rows": int(len(d)), "neurons": int((d.kind == 1).sum()), "imn": int((d.kind == 4).sum()),
         "endothelial": int((d.kind == 2).sum()), "pericytes": int((d.kind == 3).sum()), "sections": len(secs)}
    print(brain, q, flush=True)
    return q, d[d.kind == 1].subclass.value_counts()

qc = {}
a = pd.read_csv(RAW / "allen/cell_metadata_with_parcellation_annotation.csv",
                usecols=["cell_label", "brain_section_label", "class", "subclass", "x_section", "y_section", "parcellation_division"], dtype={"cell_label": str})
a = a.rename(columns={"brain_section_label": "section", "x_section": "x", "y_section": "y", "parcellation_division": "division"})
a["hq"] = True
qc["allen"], allen_counts = build(a, "allen"); del a

zcounts = {}
for i in range(1, 5):
    z = pd.read_csv(RAW / f"zhuang{i}/cell_metadata_with_cluster_annotation.csv",
                    usecols=["cell_label", "brain_section_label", "class", "subclass", "x", "y", "high_quality_transfer"], dtype={"cell_label": str})
    c = pd.read_csv(RAW / f"zhuang{i}/ccf_coordinates.csv", usecols=["cell_label", "parcellation_index"], dtype={"cell_label": str})
    z = z.merge(c, on="cell_label", how="left")
    z["division"] = z.parcellation_index.map(idx2div).fillna("unassigned")
    z = z.rename(columns={"brain_section_label": "section"}); z["hq"] = z.high_quality_transfer.astype(bool)
    h = h5py.File(RAW / f"h5ad/Zhuang-ABCA-{i}-raw.h5ad", "r")
    X = h["X"]
    if isinstance(X, h5py.Group):  # CSR
        data = X["data"][:]; ip = X["indptr"][:]
        tot = np.add.reduceat(np.r_[data, 0].astype(np.float64), ip[:-1]) * (np.diff(ip) > 0)
    else:
        tot = np.zeros(X.shape[0])
        for s0 in range(0, X.shape[0], 200000): tot[s0:s0 + 200000] = X[s0:s0 + 200000].astype(np.float64).sum(1)
    labs = h["obs"]["cell_label"]
    labs = labs[:] if isinstance(labs, h5py.Dataset) else labs["categories"][:][labs["codes"][:]]
    labs = [l.decode() if isinstance(l, bytes) else str(l) for l in labs]
    lc = pd.Series(np.log10(np.maximum(tot, 1.0)), index=labs)
    qc[f"zhuang{i}"], zc = build(z, f"zhuang{i}", lc); zcounts[f"zhuang{i}"] = zc

# ---------- frozen subclass list ----------
sub = pd.DataFrame({"allen_neurons": allen_counts})
for b, zc in zcounts.items(): sub[f"{b}_neurons"] = zc
sub = sub.fillna(0).astype(int)
sub2cls = piv.groupby("subclass")["class"].first()
sub["class"] = sub2cls.reindex(sub.index).fillna("").values
sub["n_cells_10x"] = expr.n_cells_10x.reindex(sub.index).fillna(0).astype(int)
sub["primary"] = (sub.allen_neurons >= P["min_neurons_primary"]) & (sub.n_cells_10x >= P["min_cells_10x"])
sub.index.name = "subclass"
sub.sort_index().to_csv(OUT / "subclasses.csv")
qc["n_primary_subclasses"] = int(sub.primary.sum())
json.dump(qc, open(OUT / "qc_counts.json", "w"), indent=1, sort_keys=True)
print("primary subclasses:", int(sub.primary.sum()))