#!/usr/bin/env python3
"""
Prepare MuSiC single-cell reference from Consensus-WMB-Macosko-10X.

MuSiC requires per-cell:  cell type label  +  donor/subject ID
Key difference from CIBERSORTx:
  - Uses RAW counts (not log2-normalized)
  - Sampling is balanced across donors per cell type
    (MuSiC uses cross-donor variance — every donor should be represented)

Outputs (all in --outdir):
  sc_counts.mtx.gz   sparse raw count matrix  (genes × cells)
  sc_genes.txt       gene symbols
  sc_meta.csv        cell_label, celltype, donor, donor_sex, region

Usage
-----
  python3 prep_music_ref.py
  python3 prep_music_ref.py --label_col class --max_per_donor 20
  python3 prep_music_ref.py --label_col subclass --max_per_donor 10
"""

import anndata as ad
import numpy   as np
import pandas  as pd
import scipy.sparse as sp
import scipy.io
import os, sys, argparse, gc, gzip

# ── paths ─────────────────────────────────────────────────────────────────────
BASE     = "/BRC/yan/BPA/sc_brain_ref"
EXPR_DIR = f"{BASE}/Consensus_WMB_Macosko/expression_matrices"
META_DIR = f"{BASE}/Consensus_WMB_Macosko/metadata"
TAXO_DIR = f"{BASE}/Consensus_WMB/taxonomy"
C2C_CACHE = f"{BASE}/Consensus_WMB_Macosko/cell_to_cluster_macosko_only.csv"

# ── args ──────────────────────────────────────────────────────────────────────
parser = argparse.ArgumentParser()
parser.add_argument("--label_col",     default="subclass",
                    choices=["class", "subclass", "supertype", "neurotransmitter"])
parser.add_argument("--max_per_donor", type=int, default=10,
                    help="Max cells per (cell type × donor) combination")
parser.add_argument("--min_cells",     type=int, default=10,
                    help="Drop cell types with fewer total cells")
parser.add_argument("--outdir",        default=None)
parser.add_argument("--seed",          type=int, default=42)
args = parser.parse_args()

if args.outdir is None:
    args.outdir = f"{BASE}/MuSiC_ref_{args.label_col}_maxpd{args.max_per_donor}"
os.makedirs(args.outdir, exist_ok=True)

rng = np.random.default_rng(args.seed)

# ══════════════════════════════════════════════════════════════════════════════
# STEP 1  taxonomy pivot
# ══════════════════════════════════════════════════════════════════════════════
print("[1] Building taxonomy pivot ...", flush=True)

c2ca = pd.read_csv(f"{TAXO_DIR}/cluster_to_cluster_annotation_membership.csv")
tax  = c2ca.pivot_table(
    index   = "cluster_alias",
    columns = "cluster_annotation_term_set_name",
    values  = "cluster_annotation_term_name",
    aggfunc = "first"
)
tax.index = tax.index.astype(int)

if args.label_col not in tax.columns:
    sys.exit(f"ERROR: '{args.label_col}' not available. "
             f"Options: {list(tax.columns)}")
print(f"    Using label: {args.label_col} "
      f"({tax[args.label_col].nunique()} unique values)")

# ══════════════════════════════════════════════════════════════════════════════
# STEP 2  cell → cluster → cell type
# ══════════════════════════════════════════════════════════════════════════════
print("\n[2] Loading cell → cluster mapping ...", flush=True)

if os.path.exists(C2C_CACHE):
    print(f"    Using cache: {C2C_CACHE}")
    c2c_raw = pd.read_csv(C2C_CACHE)
else:
    chunks, total_read = [], 0
    for chunk in pd.read_csv(f"{TAXO_DIR}/cell_to_cluster_membership.csv",
                             chunksize=500_000):
        total_read += len(chunk)
        hit = chunk[chunk["cell_label"].str.startswith("pBICCN")]
        if len(hit):
            chunks.append(hit)
        print(f"    read {total_read/1e6:.1f}M / 7.6M rows\r", end="", flush=True)
    print()
    c2c_raw = pd.concat(chunks)
    c2c_raw.to_csv(C2C_CACHE, index=False)
    print(f"    Cached to {C2C_CACHE}")

c2c = c2c_raw.set_index("cell_label")
c2c["cluster_alias"] = c2c["cluster_alias"].astype(int)
c2c["celltype"] = tax.loc[c2c["cluster_alias"].values, args.label_col].values

# ══════════════════════════════════════════════════════════════════════════════
# STEP 3  join donor info  (cell → library → donor)
# ══════════════════════════════════════════════════════════════════════════════
print("\n[3] Joining donor information ...", flush=True)

cell_meta = pd.read_csv(f"{META_DIR}/cell_metadata.csv",
                        usecols=["cell_label", "library_label"])
library   = pd.read_csv(f"{META_DIR}/library.csv",
                        usecols=["library_label", "donor_label",
                                 "region_of_interest_acronym"])
donor     = pd.read_csv(f"{META_DIR}/donor.csv",
                        usecols=["donor_label", "donor_sex"])

cell_meta = (cell_meta
             .merge(library, on="library_label", how="left")
             .merge(donor,   on="donor_label",   how="left")
             .set_index("cell_label"))

# join onto taxonomy table
c2c = c2c.join(cell_meta[["donor_label", "donor_sex",
                           "region_of_interest_acronym"]], how="left")

missing_donor = c2c["donor_label"].isna().sum()
if missing_donor:
    print(f"    WARNING: {missing_donor:,} cells missing donor info — dropping")
    c2c = c2c.dropna(subset=["donor_label"])

n_donors = c2c["donor_label"].nunique()
print(f"    Cells with donor info : {len(c2c):,}")
print(f"    Unique donors         : {n_donors}")
print(f"    Donor sex breakdown   : "
      f"{c2c.groupby('donor_sex').size().to_dict()}")

# ══════════════════════════════════════════════════════════════════════════════
# STEP 4  balanced subsample per (cell type × donor)
# ══════════════════════════════════════════════════════════════════════════════
print(f"\n[4] Balanced subsampling "
      f"(max {args.max_per_donor} cells per cell-type × donor) ...", flush=True)

# drop rare cell types
counts     = c2c["celltype"].value_counts()
keep_types = counts[counts >= args.min_cells].index
dropped    = counts[counts < args.min_cells].index.tolist()
if dropped:
    print(f"    Dropping {len(dropped)} types with < {args.min_cells} total cells")
c2c = c2c[c2c["celltype"].isin(keep_types)]

sampled = []
for (ct, donor), grp in c2c.groupby(["celltype", "donor_label"]):
    idx = grp.index.tolist()
    if len(idx) > args.max_per_donor:
        idx = list(rng.choice(idx, args.max_per_donor, replace=False))
    sampled.extend(idx)

c2c = c2c.loc[sampled]
sampled_set = set(c2c.index)

print(f"    Total cells selected  : {len(sampled_set):,}")
print(f"\n    Cell type summary:")
ct_summary = c2c.groupby("celltype").agg(
    n_cells  = ("donor_label", "count"),
    n_donors = ("donor_label", "nunique")
)
print(ct_summary.to_string())

# ══════════════════════════════════════════════════════════════════════════════
# STEP 5  extract RAW counts from h5ad files
# ══════════════════════════════════════════════════════════════════════════════
print(f"\n[5] Extracting RAW counts from h5ad files ...", flush=True)

# use -raw.h5ad for MuSiC (not log2)
h5ad_files = sorted([
    os.path.join(EXPR_DIR, f)
    for f in os.listdir(EXPR_DIR)
    if f.endswith("-raw.h5ad")
])

all_expr      = []
all_barcodes  = []
gene_names    = None

for h5path in h5ad_files:
    region = os.path.basename(h5path).replace("-raw.h5ad", "")
    adata  = ad.read_h5ad(h5path, backed="r")

    local_cells = [c for c in adata.obs.index if c in sampled_set]
    if not local_cells:
        adata.file.close()
        continue

    print(f"    {region}: {len(local_cells):,} cells", flush=True)

    if gene_names is None:
        if "gene_symbol" in adata.var.columns:
            gene_names = adata.var["gene_symbol"].astype(str).values
        else:
            gene_names = adata.var_names.astype(str).values

    # sort for efficient sequential HDF5 read
    idx_pos      = adata.obs.index.get_indexer(local_cells)
    sort_ord     = np.argsort(idx_pos)
    idx_sorted   = idx_pos[sort_ord]
    cells_sorted = [local_cells[i] for i in sort_ord]

    X_local = adata.X[idx_sorted, :]
    if not sp.issparse(X_local):
        X_local = sp.csr_matrix(X_local)
    else:
        X_local = X_local.tocsr()

    all_expr.append(X_local)
    all_barcodes.extend(cells_sorted)

    adata.file.close()
    del adata
    gc.collect()

# deduplicate gene symbols
_, uniq_idx = np.unique(gene_names, return_index=True)
uniq_idx.sort()
gene_names = gene_names[uniq_idx]

X_all = sp.vstack(all_expr)[:, uniq_idx]   # cells × genes  (sparse)
X_all = X_all.T.tocsc()                    # genes × cells  (sparse)

print(f"\n    Final matrix: {X_all.shape[0]:,} genes × {X_all.shape[1]:,} cells")

# ══════════════════════════════════════════════════════════════════════════════
# STEP 6  save outputs
# ══════════════════════════════════════════════════════════════════════════════
print(f"\n[6] Saving to {args.outdir} ...", flush=True)

# sparse matrix
mtx_path = f"{args.outdir}/sc_counts.mtx"
scipy.io.mmwrite(mtx_path, X_all)
import subprocess
subprocess.run(["gzip", "-f", mtx_path])
print(f"    sc_counts.mtx.gz")

# gene names
with open(f"{args.outdir}/sc_genes.txt", "w") as f:
    f.write("\n".join(gene_names))
print(f"    sc_genes.txt  ({len(gene_names):,} genes)")

# cell metadata (order matches matrix columns)
meta_out = c2c.loc[all_barcodes,
                   ["celltype", "donor_label", "donor_sex",
                    "region_of_interest_acronym"]].copy()
meta_out.index.name = "cell_label"
meta_out.columns    = ["celltype", "donor", "donor_sex", "region"]
meta_out.to_csv(f"{args.outdir}/sc_meta.csv")
print(f"    sc_meta.csv   ({len(meta_out):,} cells)")

print(f"\n{'='*60}")
print(f"Done.  Output: {args.outdir}/")
print(f"  sc_counts.mtx.gz  genes × cells  (raw counts, sparse)")
print(f"  sc_genes.txt")
print(f"  sc_meta.csv       celltype | donor | donor_sex | region")
print(f"\nNext: run run_music.R with your bulk RNA-seq data")
print(f"{'='*60}")