"""External H&E<->RNA test sets: DLPFC (spatialLIBD Visium) and breast IDC (HEST Xenium pseudo-spots).

Run:  .venv/bin/python -I 36_external_prep.py [dlpfc] [idc]     (default: both)
Needs the public spatialLIBD files in ext/external_raw/dlpfc_spatialLIBD/ (downloaded once, data only):
  <sid>_filtered_feature_bc_matrix.h5   https://spatial-dlpfc.s3.us-east-2.amazonaws.com/h5/
  <sid>_tissue_positions_list.txt       https://raw.githubusercontent.com/LieberInstitute/HumanPilot/master/10X/<sid>/
  <sid>_scalefactors_json.json          (same)

Findings this script relies on (verified, numbers printed again at run time):
DLPFC  visium_gene_expression_list.pkl = log1p(count * 3034 / total_counts_all_33538_genes)   (3034 = median
       total over all 47,681 spots, i.e. scanpy normalize_total default) restricted to the top-10000
       scanpy 'seurat'-flavour HVGs computed on all 12 samples concatenated, columns in 10x feature (genome) order.
       Rows = h5 barcode order of each sample, samples concatenated in sorted order (manifest spot_idx).
       Images = ImageNet-normalised uint8 RGB, native 224x224 full-res crop centred on the spot.
       Full-res scale: Visium centre-to-centre pitch 100 um = 138.0 px -> 0.7246 um/px (HEST: 0.725);
       note the 'pxl_row/col' columns of the HumanPilot tissue_positions_list are (x, y), not (row, col).
IDC    rebuilt from HEST-1k (MahmoodLab/hest): st/<id>.h5ad + the selected 224-px, 0.5 um/px patches of
       patches/<id>.h5 for TENX95, TENX99, NCBI783, NCBI785 (= HEST-Bench IDC), paired by barcode and checked by
       patch-centre vs spot-centre distance. Spots: in tissue, with a patch, >= 20 panel counts; > 1500 -> random
       1500 (seed 0). The earlier IDC built from ../IDC/UNI was unusable: subsample.py drew its image indices
       independently of its adata indices (key '/k/__list__' vs '/k/__arr__', hash()-randomised seed).
       First run fetches patches with HTTP range reads: needs HF_TOKEN in the environment.
Outputs cache/external/: see the OUT dict printed at the end.
"""
import json
import os
import pickle
import sys

import h5py
import numpy as np
import pandas as pd
import scipy.sparse as sp
from PIL import Image

ROOT = os.path.dirname(os.path.abspath(__file__))
DATA = os.path.abspath(os.path.join(ROOT, ".."))
RAW = os.path.join(ROOT, "ext", "external_raw", "dlpfc_spatialLIBD")
OUT = os.path.join(ROOT, "cache", "external")
os.makedirs(OUT, exist_ok=True)
IMNET_MEAN = np.array([0.485, 0.456, 0.406], np.float32)
IMNET_STD = np.array([0.229, 0.224, 0.225], np.float32)
TARGET_UM = 0.5  # um/px of the resampled patches
HVG = pd.read_csv(os.path.join(ROOT, "cache", "hest", "hvg_genes.csv"))["gene"].astype(str).tolist()
HEST_META = pd.read_csv(os.path.join(DATA, "HEST", "HEST_v1_1_0.csv"))
HEST_TRAIN = set(pd.read_csv(os.path.join(ROOT, "cache", "hest", "meta.csv"), usecols=["sample"])["sample"])


def load_pickle(p):
    with open(p, "rb") as f:
        return pickle.load(f)


def denorm(x):
    """(N,3,H,W) ImageNet-normalised float -> (N,H,W,3) uint8; returns max dequantisation error (in 0..255 units)."""
    v = (np.asarray(x, np.float32).transpose(0, 2, 3, 1) * IMNET_STD + IMNET_MEAN) * 255.0
    u = np.clip(np.round(v), 0, 255).astype(np.uint8)
    return u, float(np.abs(v - u).max())


def resample(u8, um_native, fov_um=224 * TARGET_UM, out=224):
    """Centre crop fov_um (if available) and resize to out x out => TARGET_UM um/px."""
    res = []
    for im in u8:
        side = fov_um / um_native
        h = im.shape[0]
        if side > h:  # not enough context: keep full patch (FOV smaller than requested)
            side = h
        c = (h - side) / 2
        res.append(np.asarray(Image.fromarray(im).resize((out, out), Image.BICUBIC, box=(c, c, c + side, c + side))))
    return np.stack(res)


def hvg_overlap(genes):
    g = set(genes)
    return int(sum(x in g for x in HVG))


# ================================================================== DLPFC
def read_h5(sid):
    with h5py.File(os.path.join(RAW, f"{sid}_filtered_feature_bc_matrix.h5"), "r") as f:
        g = f["matrix"]
        M = sp.csc_matrix((g["data"][:], g["indices"][:], g["indptr"][:]), shape=g["shape"][:]).T.tocsr()
        bcs = [b.decode() for b in g["barcodes"][:]]
        names = [b.decode() for b in g["features"]["name"][:]]
        ids = [b.decode() for b in g["features"]["id"][:]]
    return M, bcs, names, ids


def seurat_hvg(L, n_top):
    """scanpy.pp.highly_variable_genes(flavor='seurat', n_top_genes=n_top) on log1p data L (csr)."""
    E = L.copy()
    E.data = np.expm1(E.data)
    n = E.shape[0]
    mean = np.asarray(E.mean(0)).ravel()
    var = (np.asarray(E.multiply(E).mean(0)).ravel() - mean ** 2) * n / (n - 1)
    with np.errstate(all="ignore"):
        disp = var / mean
        disp[disp == 0] = np.nan
        disp = np.log(disp)
    df = pd.DataFrame({"m": np.log1p(mean), "d": disp})
    df["bin"] = pd.cut(df.m, bins=20)
    g = df.groupby("bin", observed=False)["d"]
    bm, bs, one = g.transform("mean"), g.transform("std"), g.transform("count") == 1
    bs, bm = bs.where(~one, bm), bm.where(~one, 0)
    dn = ((df.d - bm) / bs).to_numpy()
    v = np.sort(dn[~np.isnan(dn)])[::-1]
    cut = v[n_top - 1]
    print(f"  HVG cut-off normalised dispersion {cut:.5f}; next {v[n_top]:.5f} (gap {cut - v[n_top]:.2e})")
    return np.where(np.nan_to_num(dn) >= cut)[0]


def dlpfc():
    print("== DLPFC")
    P = os.path.join(DATA, "DLPFC_visium", "GPFM")
    man = json.load(open(os.path.join(DATA, "_manifest", "dlpfc.json")))
    samples = man["samples"]
    X = load_pickle(os.path.join(P, "visium_gene_expression_list.pkl"))
    X = np.asarray(X.numpy() if hasattr(X, "numpy") else X, np.float64)
    gidx = np.asarray(pd.read_csv(os.path.join(DATA, "DLPFC_visium", "global_spot_index.csv"))["global_spot_idx"])
    assert np.array_equal(gidx, man["global_spot_idx"]) and X.shape == (len(gidx), 10000)

    Ms, rows, names, ids = [], [], None, None
    sf_by, pos_by, bcs_by = {}, {}, {}
    for sid in samples:
        M, bcs, names, ids = read_h5(sid)
        assert M.shape[0] == man["n_spots"][sid], sid
        Ms.append(M)
        bcs_by[sid] = bcs
        sf_by[sid] = json.load(open(os.path.join(RAW, f"{sid}_scalefactors_json.json")))
        pos = pd.read_csv(os.path.join(RAW, f"{sid}_tissue_positions_list.txt"), header=None,
                          names=["barcode", "in_tissue", "array_row", "array_col", "px_x", "px_y"]).set_index("barcode")
        pos_by[sid] = pos
        for i in man["spot_idx"][sid]:
            rows.append((sid, i, bcs[i]))
    M = sp.vstack(Ms).tocsr().astype(np.float64)
    tot = np.asarray(M.sum(1)).ravel()
    target = float(np.median(tot))
    A = M[gidx]  # raw counts of the 482 spots, all 33538 genes
    print(f"  all spots {M.shape}, median total counts {target}")

    # 1) normalisation: x = log1p(count * target / total)
    L = sp.diags(target / tot) @ M
    L.data = np.log1p(L.data)
    # 2) gene subset: reproduce scanpy seurat HVG (top 10000) on all 47,681 spots
    sel = seurat_hvg(L, 10000)
    assert len(sel) == 10000
    Lsel = L[gidx][:, sel].toarray()
    err = np.abs(Lsel - X).max()
    # 3) independent check: exact per-column value match of recovered counts vs every h5 gene
    Ad = A.toarray()
    K = np.round(np.expm1(X) * (tot[gidx] / target)[:, None])
    buckets = {}
    for j in range(Ad.shape[1]):
        buckets.setdefault(Ad[:, j].tobytes(), []).append(j)
    cand = [buckets.get(K[:, c].tobytes(), []) for c in range(10000)]
    n_unique = sum(len(c) == 1 for c in cand)
    n_consistent = sum(sel[c] in cand[c] for c in range(10000))
    n_unique_ok = sum(len(cand[c]) == 1 and cand[c][0] == sel[c] for c in range(10000))
    print(f"  max |log-expr(reconstructed) - pickle| = {err:.2e}")
    print(f"  exact-value column matches: {n_unique}/10000 columns match exactly one h5 gene "
          f"({n_unique_ok} of them = HVG reconstruction); {n_consistent}/10000 HVG genes value-consistent; "
          f"{int((K.sum(0) == 0).sum())} columns all-zero in the 482 spots")
    assert err < 1e-5 and n_consistent == 10000

    genes = pd.DataFrame({"col": np.arange(10000), "gene": np.array(names)[sel], "ensembl_id": np.array(ids)[sel],
                          "h5_feature_idx": sel,
                          "match": ["exact_unique" if len(c) == 1 else "exact_ambiguous_resolved_by_hvg" for c in cand]})
    genes["duplicated_symbol"] = genes.gene.duplicated(keep=False)
    genes.to_csv(os.path.join(OUT, "dlpfc_genes.csv"), index=False)
    np.save(os.path.join(OUT, "dlpfc_counts.npy"), Ad[:, sel].astype(np.float32))
    # full transcriptome version (same spots) from the public h5
    np.save(os.path.join(OUT, "dlpfc_counts_full.npy"), Ad.astype(np.float32))
    pd.DataFrame({"gene": names, "ensembl_id": ids, "in_10k": np.isin(np.arange(len(names)), sel)}).to_csv(
        os.path.join(OUT, "dlpfc_genes_full.csv"), index=False)

    # images
    im = load_pickle(os.path.join(P, "visium_image_list.pkl"))
    im = im.numpy() if hasattr(im, "numpy") else np.asarray(im)
    u8, qerr = denorm(im)
    print(f"  images {im.shape}: ImageNet-normalised, dequantisation error {qerr:.1e} (exact uint8)")
    np.save(os.path.join(OUT, "dlpfc_he.npy"), u8)

    hest_ids = dict(zip(HEST_META.subseries.astype(str), HEST_META.id))
    meta = []
    for k, (sid, i, bc) in enumerate(rows):
        p = pos_by[sid]
        q = p.reset_index()
        nb = q.merge(q, on="array_row")
        nb = nb[nb.array_col_y == nb.array_col_x + 2]
        pitch = float(np.median(np.hypot(nb.px_x_x - nb.px_x_y, nb.px_y_x - nb.px_y_y)))  # 100 um
        hid = hest_ids.get(sid)
        meta.append(dict(sample=sid, spot_idx=i, global_spot_idx=int(gidx[k]), barcode=bc,
                         array_row=int(p.loc[bc, "array_row"]), array_col=int(p.loc[bc, "array_col"]),
                         x_fullres=float(p.loc[bc, "px_x"]), y_fullres=float(p.loc[bc, "px_y"]),
                         organ="Brain", tissue="dorsolateral prefrontal cortex",
                         um_per_px=100.0 / pitch, fov_um=224 * 100.0 / pitch,
                         spot_diameter_fullres=sf_by[sid]["spot_diameter_fullres"],
                         total_counts=float(tot[gidx[k]]), hest_id=hid, in_hest_train_cache=hid in HEST_TRAIN))
    meta = pd.DataFrame(meta)
    meta.to_csv(os.path.join(OUT, "dlpfc_meta.csv"), index=False)
    um = float(meta.um_per_px.median())
    u05 = resample(u8, um)
    np.save(os.path.join(OUT, "dlpfc_he_05um.npy"), u05)
    print(f"  um/px {um:.4f} -> patch FOV {224 * um:.1f} um; dlpfc_he_05um = centre {112 / um:.1f}px crop -> 224 (112 um)")
    print(f"  HEST-HVG overlap: 10k cols {hvg_overlap(genes.gene)}/{len(HVG)}, full {hvg_overlap(names)}/{len(HVG)}")


# ================================================================== IDC (paired, from HEST-1k)
HEST_RAW = os.path.join(ROOT, "ext", "external_raw", "hest_idc")
IDC_IDS = ["TENX95", "TENX99", "NCBI783", "NCBI785"]  # = HEST-Bench IDC
IDC_MIN_COUNTS = 20
IDC_MAX_SPOTS = 1500
CTRL = ("BLANK", "NegControl", "Unassigned", "antisense", "DeprecatedCodeword", "Intergenic")


def fetch_patches(sid, barcodes):
    """Fetch only the selected patches of MahmoodLab/hest patches/<sid>.h5 by HTTP range reads (the full files are
    5.0 GB together; patches are stored uncompressed, one chunk per patch). Auth: HF_TOKEN env var, read by
    huggingface_hub itself. Writes patches_subset/<sid>.h5 with img/barcode/coords + the source attrs."""
    from huggingface_hub import HfFileSystem
    out = os.path.join(HEST_RAW, "patches_subset", f"{sid}.h5")
    os.makedirs(os.path.dirname(out), exist_ok=True)
    fs = HfFileSystem()
    src = f"datasets/MahmoodLab/hest/patches/{sid}.h5"
    pb = all_patch_barcodes(sid)  # cached barcode/coords table of the remote file (row order = h5 row order)
    coords = pb[["x", "y"]].to_numpy()
    pos = {b: i for i, b in enumerate(pb.barcode)}
    rows = np.array([pos[b] for b in barcodes])
    order = np.argsort(rows)
    imgs = np.zeros((len(rows), 224, 224, 3), np.uint8)
    with fs.open(src, "rb", block_size=160 * 1024, cache_type="readahead") as f, h5py.File(f, "r") as h:
        attrs = dict(h["img"].attrs)
        for n, k in enumerate(order):
            imgs[k] = h["img"][rows[k]]
            if n % 250 == 0:
                print(f"    {sid}: fetched {n}/{len(rows)} patches", flush=True)
    with h5py.File(out, "w") as g:
        g["img"] = imgs
        g["barcode"] = np.array(barcodes, dtype="S")
        g["coords"] = coords[rows]
        g["src_row"] = rows
        for k, v in attrs.items():
            g["img"].attrs[k] = v
    return out


def all_patch_barcodes(sid):
    """Barcodes that have a patch in HEST (read remotely once, cached as a small csv)."""
    p = os.path.join(HEST_RAW, "patches_subset", f"{sid}_all_barcodes.csv")
    if not os.path.exists(p):
        from huggingface_hub import HfFileSystem
        os.makedirs(os.path.dirname(p), exist_ok=True)
        with HfFileSystem().open(f"datasets/MahmoodLab/hest/patches/{sid}.h5", "rb", block_size=2 ** 22) as f, \
                h5py.File(f, "r") as h:
            bc = [b[0].decode() if isinstance(b[0], bytes) else str(b[0]) for b in h["barcode"][:]]
            pd.DataFrame({"barcode": bc, "x": h["coords"][:, 0], "y": h["coords"][:, 1]}).to_csv(p, index=False)
    return pd.read_csv(p)


def idc():
    print("== IDC (HEST-1k, paired by barcode)")
    import anndata as ad
    ads = {s: ad.read_h5ad(os.path.join(HEST_RAW, "st", f"{s}.h5ad")) for s in IDC_IDS}
    common = set.intersection(*[set(a.var_names) for a in ads.values()])
    genes = sorted(g for g in common if not g.startswith(CTRL))
    rng = np.random.default_rng(0)
    he, counts, meta = [], [], []
    for sid in IDC_IDS:
        a = ads[sid]
        md = json.load(open(os.path.join(HEST_RAW, "metadata", f"{sid}.json")))
        X = a[:, genes].X
        X = np.asarray(X.toarray() if sp.issparse(X) else X, np.float32)
        tot = X.sum(1)
        pb = all_patch_barcodes(sid)
        in_t = a.obs["in_tissue"].astype(bool).values if "in_tissue" in a.obs else np.ones(a.n_obs, bool)
        has_p = np.isin(a.obs_names, pb.barcode)
        elig = np.where(in_t & has_p & (tot >= IDC_MIN_COUNTS))[0]
        sub = len(elig) > IDC_MAX_SPOTS
        sel = np.sort(rng.choice(elig, IDC_MAX_SPOTS, replace=False)) if sub else elig
        print(f"  {sid}: spots {a.n_obs}, in_tissue {in_t.sum()}, with patch {has_p.sum()}, "
              f">= {IDC_MIN_COUNTS} counts {len(elig)} -> kept {len(sel)}{' (random subsample, seed 0)' if sub else ''}")
        bcs = a.obs_names[sel].tolist()
        p = os.path.join(HEST_RAW, "patches_subset", f"{sid}.h5")
        if not os.path.exists(p):
            fetch_patches(sid, bcs)
        with h5py.File(p, "r") as g:
            assert [b.decode() for b in g["barcode"][:]] == bcs, f"{sid}: cached patch subset does not match selection"
            imgs, coords = g["img"][:], g["coords"][:]
            src_px, ds, mpp = (float(g["img"].attrs[k]) for k in ("patch_size_src", "downsample", "pixel_size"))
        # pairing check: patch centre (top-left coords + half the source patch size) vs spot centre
        cx, cy = coords[:, 0] + src_px / 2, coords[:, 1] + src_px / 2
        o = a.obs.iloc[sel]
        dev = np.hypot(cx - o.pxl_col_in_fullres.values, cy - o.pxl_row_in_fullres.values) * mpp
        um_out = mpp * ds
        print(f"    patch: {int(src_px)} src px @ {mpp:.4f} um/px -> 224 px @ {um_out:.4f} um/px ({224 * um_out:.0f} um FOV); "
              f"|patch centre - spot centre| max {dev.max():.2f} um, median {np.median(dev):.2f} um")
        assert dev.max() < 5, f"{sid}: patch/spot coordinate mismatch"
        he.append(imgs)
        counts.append(X[sel])
        for j, i in enumerate(sel):
            meta.append(dict(sample=sid, barcode=a.obs_names[i], array_row=int(a.obs.array_row.iloc[i]),
                             array_col=int(a.obs.array_col.iloc[i]),
                             x_fullres=float(a.obs.pxl_col_in_fullres.iloc[i]),
                             y_fullres=float(a.obs.pxl_row_in_fullres.iloc[i]),
                             organ="Breast", oncotree=md.get("oncotree_code"), um_per_px=um_out,
                             um_per_px_fullres=mpp, patch_center_dev_um=float(dev[j]), total_counts=float(tot[i]),
                             subsampled=bool(sub), n_eligible=len(elig),
                             in_hest_train_cache=sid in HEST_TRAIN))
    he, counts = np.concatenate(he), np.concatenate(counts)
    np.save(os.path.join(OUT, "idc_he.npy"), he)
    np.save(os.path.join(OUT, "idc_counts.npy"), counts)
    pd.DataFrame({"gene": genes}).to_csv(os.path.join(OUT, "idc_genes.csv"), index=False)
    pd.DataFrame(meta).to_csv(os.path.join(OUT, "idc_meta.csv"), index=False)
    for f in ("idc_he_unpaired.npy", "idc_he_unpaired_meta.csv"):  # old, unpaired subsample images: removed
        if os.path.exists(os.path.join(OUT, f)):
            os.remove(os.path.join(OUT, f))
    hest_common = pd.read_csv(os.path.join(ROOT, "cache", "hest", "genes.csv"))["gene"].astype(str)
    print(f"  genes {len(genes)} (common non-control panel genes); counts {counts.shape}; images {he.shape}")
    print(f"  IDC panel overlap: HEST HVG {hvg_overlap(genes)}/{len(HVG)}; "
          f"HEST common genes {int(hest_common.isin(genes).sum())}/{len(hest_common)}")


if __name__ == "__main__":
    todo = sys.argv[1:] or ["dlpfc", "idc"]
    if "dlpfc" in todo:
        dlpfc()
    if "idc" in todo:
        idc()
    for f in sorted(os.listdir(OUT)):
        p = os.path.join(OUT, f)
        print(f"  {f:28s}", np.load(p, mmap_mode="r").shape if f.endswith(".npy") else f"{len(pd.read_csv(p))} rows")