"""Prepare one split before fitting expression features and the PCA basis.

The cache preserves cell IDs, feature order, split assignments, input checksums,
and the fitted transformation. All operations that estimate parameters use
training cells. Library-size normalization acts independently on each cell.
"""

from __future__ import annotations
import hashlib, json
from pathlib import Path
from importlib.metadata import version
import numpy as np
import pandas as pd
import scipy.sparse as sp
from sklearn.decomposition import PCA

DATASETS = {
    "norman": {"file": "NormanWeissman2019_filtered.h5ad", "operation": "activation"},
    "replogle_k562": {
        "file": "ReplogleWeissman2022_K562_essential.h5ad",
        "operation": "interference",
    },
}


def checksum(path: str | Path) -> str:
    """SHA-256 of file bytes, read in bounded blocks."""
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for b in iter(lambda: f.read(8 * 1024**2), b""):
            h.update(b)
    return h.hexdigest()


def parse(label: str, control: str = "control", sep: str = "_") -> list[str]:
    """Return a sorted gene set, excluding the explicit control token."""
    genes = [g for g in str(label).split(sep) if g and g != control]
    if len(genes) != len(set(genes)):
        raise ValueError(f"Duplicate gene in perturbation {label!r}")
    return sorted(genes)


def assign_splits(
    obs: pd.DataFrame, regime: str, seed: int, control="control", sep="_"
) -> np.ndarray:
    """Assign train/validation/test using cell, label, gene, or combination units.

    Target labels are split 70/10/20 for perturbation and combination regimes.
    Each cell group and the controls are split 70/15/15 in the cell regime.
    Gene holdouts exclude every combination containing a held-out gene.
    """
    rng = np.random.default_rng(seed)
    labels = sorted(set(obs.perturbation) - {control})
    split = np.full(len(obs), "train", dtype="U5")
    ctrl = np.flatnonzero(obs.is_control.to_numpy())
    if len(ctrl) < 10:
        raise ValueError("At least 10 control cells are required")

    def divide(ids):
        ids = rng.permutation(ids)
        n = len(ids)
        if n < 5:
            raise ValueError("At least five observations per split unit are required")
        a, b = int(0.7 * n), int(0.85 * n)
        split[ids[a:b]] = "val"
        split[ids[b:]] = "test"

    divide(ctrl)
    if regime == "cell":
        for label in labels:
            divide(np.flatnonzero(obs.perturbation.to_numpy() == label))
    else:
        units = labels
        if regime == "combination":
            units = [p for p in labels if len(parse(p, control, sep)) == 2]
        elif regime == "gene":
            units = sorted({g for p in labels for g in parse(p, control, sep)})
        elif regime != "perturbation":
            raise ValueError(f"Unsupported split {regime}")
        if len(units) < 5:
            raise ValueError(f"{regime} needs at least five independent units")
        perm = rng.permutation(units)
        nt = max(1, int(0.2 * len(units)))
        nv = max(1, int(0.1 * len(units)))
        test, val = set(perm[:nt]), set(perm[nt : nt + nv])
        for p in labels:
            genes = set(parse(p, control, sep))
            s = (
                "test"
                if (genes & test if regime == "gene" else p in test)
                else (
                    "val"
                    if (genes & val if regime == "gene" else p in val)
                    else "train"
                )
            )
            split[obs.perturbation.to_numpy() == p] = s
    if not all(
        np.any((split == s) & ~obs.is_control.to_numpy())
        for s in ["train", "val", "test"]
    ):
        raise ValueError("Every partition requires perturbed cells")
    return split


def prepare(
    raw: str,
    output: str,
    dataset: str,
    regime="perturbation",
    seed=0,
    input_scale="counts",
    n_hvg=2000,
    n_pca=50,
    min_cells=20,
    max_cells=None,
    max_per_group=None,
    max_groups=None,
    pert_col="perturbation",
    batch_col="gemgroup",
    celltype_col="celltype",
    control="control",
    sep="_",
) -> dict:
    """Read h5ad, split cells, fit training HVGs/PCA, and save a self-contained cache.

    `input_scale` is explicit: counts receive CP10k and log1p; log1p values are
    used directly. The optional group cap selects labels without examining
    expression. Backed input avoids loading a whole atlas for a small run.
    """
    import anndata as ad
    import scanpy as sc

    out = Path(output)
    if (out / "meta.json").exists():
        raise FileExistsError(
            f"Cache already exists: {out}; choose a new output directory"
        )
    a = ad.read_h5ad(raw, backed="r")
    if pert_col not in a.obs:
        raise KeyError(f"Missing {pert_col}; available columns: {list(a.obs.columns)}")
    labels = np.asarray(
        [
            sep.join(parse(p, control, sep)) or control
            for p in a.obs[pert_col].astype(str)
        ]
    )
    rng = np.random.default_rng(seed)
    keep = np.arange(a.n_obs)
    if max_groups:
        groups = sorted(set(labels) - {control})
        selected = rng.choice(groups, min(max_groups, len(groups)), replace=False)
        keep = keep[np.isin(labels, list(selected) + [control])]
    if max_cells and len(keep) > max_cells:
        keep = np.sort(rng.choice(keep, max_cells, replace=False))
    if max_per_group:
        keep = np.sort(
            np.concatenate(
                [
                    rng.choice(ids, min(max_per_group, len(ids)), replace=False)
                    for p in sorted(set(labels[keep]))
                    if len(ids := keep[labels[keep] == p])
                ]
            )
        )
    vc = pd.Series(labels[keep]).value_counts()
    permitted = set(vc[vc >= min_cells].index) | {control}
    keep = keep[np.isin(labels[keep], list(permitted))]
    b = a[keep].to_memory()
    a.file.close()
    if not b.obs_names.is_unique or not b.var_names.is_unique:
        raise ValueError("Cell and feature identifiers must be unique")
    obs = pd.DataFrame(
        {
            "cell_id": b.obs_names.astype(str),
            "perturbation": labels[keep],
            "batch": (
                b.obs[batch_col].astype(str).to_numpy() if batch_col in b.obs else "0"
            ),
            "celltype": (
                b.obs[celltype_col].astype(str).to_numpy()
                if celltype_col in b.obs
                else dataset
            ),
        }
    )
    obs["is_control"] = obs.perturbation.eq(control)
    obs["split"] = assign_splits(obs, regime, seed, control, sep)
    # Reserve outcome cells for every candidate before fitting any features.
    # These cells allow evaluation of a nominated action with measured responses
    # independent of the target query and every model-training outcome.
    for label in sorted(set(obs.perturbation)):
        ids = obs.index[obs.perturbation.eq(label)].to_numpy()
        n_ref = max(2, int(0.1 * len(ids)))
        # Preserve all three control partitions when selecting reference cells.
        if label == control:
            ids = ids[obs.loc[ids, "split"].eq("train").to_numpy()]
        ref = rng.choice(ids, min(n_ref, max(0, len(ids) - 3)), replace=False)
        obs.loc[ref, "split"] = "reference"
    train = obs.split.eq("train").to_numpy()
    X = sp.csr_matrix(b.X, dtype=np.float32)
    if not np.isfinite(X.data).all() or (X.data < 0).any():
        raise ValueError("Expression must be finite and nonnegative")
    if input_scale == "counts":
        if not np.allclose(X.data, np.round(X.data), atol=1e-5):
            raise ValueError(
                "Counts mode requires integer counts; specify log1p for normalized data"
            )
        totals = np.asarray(X.sum(axis=1)).ravel()
        if np.any(totals <= 0):
            raise ValueError("Cells with zero total counts are unsupported")
        X = sp.diags(1e4 / totals) @ X
        X = X.tocsr()
        X.data = np.log1p(X.data)
    elif input_scale != "log1p":
        raise ValueError("input_scale must be counts or log1p")
    # Feature statistics are estimated without validation or test expression.
    ta = ad.AnnData(X[train].copy(), var=b.var.copy())
    if n_hvg < X.shape[1]:
        sc.pp.highly_variable_genes(ta, n_top_genes=n_hvg, flavor="seurat")
        hv = np.flatnonzero(ta.var.highly_variable.to_numpy())
    else:
        hv = np.arange(X.shape[1])
    X = X[:, hv].tocsr().astype(np.float32)
    d = min(n_pca, X.shape[1] - 1, int(train.sum()) - 1)
    if d < 1:
        raise ValueError("Insufficient dimensions for PCA")
    # ARPACK centers sparse input internally and avoids a dense full-atlas copy.
    pca = PCA(n_components=d, svd_solver="arpack", random_state=seed).fit(X[train])
    emb = pca.transform(X).astype(np.float32)
    out.mkdir(parents=True, exist_ok=True)
    sp.save_npz(out / "Xhvg.npz", X)
    np.save(out / "pca_emb.npy", emb)
    np.save(out / "pca_components.npy", pca.components_.astype(np.float32))
    np.save(out / "pca_mean.npy", pca.mean_.astype(np.float32))
    obs.to_parquet(out / "obs.parquet", index=False)
    (out / "genes_hvg.txt").write_text("\n".join(b.var_names[hv].astype(str)))
    # A fixed evaluation bandwidth supports comparisons across predicted populations.
    tr_emb = emb[train]
    samp = tr_emb[rng.choice(len(tr_emb), min(512, len(tr_emb)), replace=False)]
    from scipy.spatial.distance import pdist

    bandwidth = float(np.median(pdist(samp) ** 2))
    gamma = 1 / max(bandwidth, 1e-8)
    meta = {
        "reference_fraction": 0.1,
        "protocol": "split-first-v1",
        "dataset": dataset,
        "name": dataset,
        "operation": DATASETS.get(dataset, {}).get("operation", "interference"),
        "control_label": control,
        "sep": sep,
        "split": regime,
        "seed": seed,
        "input_scale": input_scale,
        "raw_sha256": checksum(raw),
        "raw_file": Path(raw).name,
        "n_cells": len(obs),
        "n_control": int(obs.is_control.sum()),
        "n_hvg": len(hv),
        "n_pca": d,
        "pca_explained_var": float(pca.explained_variance_ratio_.sum()),
        "mmd_gamma": gamma,
        "fit_cell_ids_sha256": hashlib.sha256(
            "\n".join(obs.loc[train, "cell_id"]).encode()
        ).hexdigest(),
        "selection": {
            "max_cells": max_cells,
            "max_per_group": max_per_group,
            "max_groups": max_groups,
        },
        "software": {"scanpy": version("scanpy"), "anndata": version("anndata")},
    }
    (out / "meta.json").write_text(json.dumps(meta, indent=2))
    return meta