"""Expression cache with explicit partitions and a metadata-only gene vocabulary."""

from __future__ import annotations
import json
from pathlib import Path
import numpy as np
import pandas as pd
import scipy.sparse as sp
from .preprocess import parse


class PerturbData:
    """Load a prepared cache. Arrays retain the row order of obs.parquet."""

    def __init__(self, cache_dir: str, embedding="pca"):
        self.dir = str(cache_dir)
        p = Path(cache_dir)
        self.meta = json.loads((p / "meta.json").read_text())
        if self.meta.get("protocol") != "split-first-v1":
            raise ValueError(
                "Use a split-first-v1 cache. Historical checkpoints require the archived code."
            )
        if embedding != "pca":
            raise ValueError(
                "This release implements the manuscript PCA representation"
            )
        self.obs = pd.read_parquet(p / "obs.parquet")
        self.genes = (p / "genes_hvg.txt").read_text().splitlines()
        self.Xhvg = sp.load_npz(p / "Xhvg.npz").tocsr()
        self.emb = np.load(p / "pca_emb.npy")
        self.d = self.emb.shape[1]
        self.pca_components = np.load(p / "pca_components.npy")
        self.pca_mean = np.load(p / "pca_mean.npy")
        if len(self.obs) != len(self.emb) or self.Xhvg.shape != (
            len(self.obs),
            len(self.genes),
        ):
            raise ValueError("Cache row or feature dimensions disagree")
        self.sep = self.meta["sep"]
        self.control_label = self.meta["control_label"]
        self.operation = self.meta["operation"]
        self.is_control = self.obs.is_control.to_numpy()
        self.control_idx = np.flatnonzero(self.is_control)
        self.batch = self.obs.batch.to_numpy()
        self.celltype = self.obs.celltype.to_numpy()
        self.pert_to_idx = {
            p: s.index.to_numpy()
            for p, s in self.obs.groupby("perturbation")
            if p != self.control_label
        }
        self.perturbations = sorted(self.pert_to_idx)
        self.genes_vocab = sorted(
            {g for p in self.perturbations for g in self.parse(p)}
        )
        self.gene_to_id = {g: i for i, g in enumerate(self.genes_vocab)}
        self.op_vocab = ["none", self.operation]
        self.op_to_id = {o: i for i, o in enumerate(self.op_vocab)}
        self.singles = [p for p in self.perturbations if len(self.parse(p)) == 1]
        self.combos = [p for p in self.perturbations if len(self.parse(p)) == 2]

    def parse(self, label: str) -> list[str]:
        return parse(label, self.control_label, self.sep)

    def indices(self, partition: str, controls: bool | None = None) -> np.ndarray:
        mask = self.obs.split.eq(partition).to_numpy(copy=True)
        if controls is not None:
            mask &= self.is_control if controls else ~self.is_control
        return np.flatnonzero(mask)

    def labels(self, partition: str) -> list[str]:
        return sorted(set(self.obs.iloc[self.indices(partition, False)].perturbation))

    def decode_to_genes(self, emb: np.ndarray) -> np.ndarray:
        """Rank-d affine reconstruction, shape (..., d) to (..., n_hvg)."""
        return emb @ self.pca_components + self.pca_mean

    def pert_gene_op_ids(self, label):
        genes = self.parse(label)
        unknown = set(genes) - set(self.gene_to_id)
        if unknown:
            raise ValueError(f"Genes outside checkpoint vocabulary: {sorted(unknown)}")
        ids = np.array([self.gene_to_id[g] for g in genes], dtype=np.int64)
        return ids, np.full(len(ids), 1, dtype=np.int64)

    def sample_controls(self, target_idx, strategy, rng, control_pool):
        """Sample only supplied controls, optionally within batch/cell type.

        Missing matched controls raise an error so cross-context substitutions
        cannot change the conditioning distribution without an explicit choice.
        """
        pool = np.asarray(control_pool, dtype=int)
        if not len(pool) or not self.is_control[pool].all():
            raise ValueError("Invalid control pool")
        if strategy == "random":
            return rng.choice(pool, len(target_idx), replace=True)
        if strategy == "nearest":
            from sklearn.neighbors import NearestNeighbors

            j = (
                NearestNeighbors(n_neighbors=1)
                .fit(self.emb[pool])
                .kneighbors(self.emb[target_idx], return_distance=False)
            )
            return pool[j.ravel()]
        if strategy not in ("batch", "celltype", "batch_celltype"):
            raise ValueError(strategy)
        out = []
        for i in target_idx:
            m = np.ones(len(pool), bool)
            if strategy in ("batch", "batch_celltype"):
                m &= self.batch[pool] == self.batch[i]
            if strategy in ("celltype", "batch_celltype"):
                m &= self.celltype[pool] == self.celltype[i]
            if not m.any():
                raise ValueError(
                    f"No matching control for cell {self.obs.iloc[i].cell_id}"
                )
            out.append(rng.choice(pool[m]))
        return np.asarray(out, dtype=np.int64)