from pathlib import Path

import scanpy as sc
import anndata as ad
import numpy as np


def load_and_qc(
    h5ad_path: Path,
    min_genes_per_cell: int = 500,
    max_genes_per_cell: int = 8000,
    max_pct_mito: float = 15.0,
    max_pct_ribo: float = 50.0,
) -> ad.AnnData:
    """Standard QC for K562 dataset."""
    adata = sc.read_h5ad(h5ad_path)
    print(f"Loading: cells {adata.n_obs}, genes {adata.n_vars}")
    
    # Mitochondria & Ribosome gene tags
    adata.var["mt"] = adata.var_names.str.startswith("MT-")
    adata.var["ribo"] = adata.var_names.str.startswith(("RPS", "RPL"))
    sc.pp.calculate_qc_metrics(adata, qc_vars=["mt", "ribo"], inplace=True)
    
    # Filtering
    sc.pp.filter_cells(adata, min_genes=min_genes_per_cell)
    sc.pp.filter_cells(adata, max_genes=max_genes_per_cell)
    adata = adata[adata.obs["pct_counts_mt"] < max_pct_mito, :].copy()
    adata = adata[adata.obs["pct_counts_ribo"] < max_pct_ribo, :].copy()
    
    # Normalization & log transformation
    sc.pp.normalize_total(adata, target_sum=1e4)
    sc.pp.log1p(adata)
    
    # HVG & scale (Geneformer uses raw counts, so keep a separate layer)
    if "raw_counts" not in adata.layers:
        adata.layers["raw_counts"] = adata.X.copy()  # For Geneformer input
    
    print(f"After QC: cells {adata.n_obs}, genes {adata.n_vars}")
    return adata

import torch


class GeneformerEmbedder:
    """Wrapper for Geneformer 95M pre-trained model.
    
    Uses EmbExtractor and TranscriptomeTokenizer from the official Geneformer repository.
    """

    MODEL_VARIANTS = {
        "gf-6L-30M-i2048": "Geneformer 30M small, context 2048",
        "gf-12L-95M-i2048": "Geneformer 95M standard, context 2048",
        "gf-12L-95M-i4096": "Geneformer 95M extended, context 4096",
    }

    def __init__(
        self,
        device: str = "cuda",
        model_variant: str = "gf-12L-95M-i2048",
    ):
        # Geneformer has its own tokenizer + BertForMaskedLM structure
        # For practical code, refer to the helical or official repository examples.
        from geneformer import TranscriptomeTokenizer, EmbExtractor
        self.tokenizer = TranscriptomeTokenizer(
            custom_attr_name_dict={"cell_type": "cell_type", "target_gene": "target_gene"},
            nproc=4,
        )
        self.extractor = EmbExtractor(
            model_type="Pretrained",
            num_classes=0,
            emb_mode="cell",       # Cell embedding (similar to CLS token)
            filter_data=None,
            max_ncells=None,
            emb_layer=-1,
            forward_batch_size=8,
            nproc=4,
        )
        self.device = device
        self.model_variant = model_variant

    def tokenize_adata(self, adata: ad.AnnData, output_dir: Path) -> Path:
        """Convert AnnData to Geneformer token file (.dataset).
        
        Requires Ensembl gene ID (adata.var["ensembl_id"]).
        """
        if "ensembl_id" not in adata.var.columns:
            # Map symbols to Ensembl IDs using mygene or pyensembl
            raise ValueError("adata.var['ensembl_id'] required. Map using mygene or pyensembl.")
        
        output_dir.mkdir(parents=True, exist_ok=True)
        # Convert anndata to loom format and then tokenize (official Geneformer workflow)
        loom_path = output_dir / "cells.loom"
        adata.write_loom(loom_path)
        self.tokenizer.tokenize_data(
            data_directory=str(output_dir),
            output_directory=str(output_dir),
            output_prefix="tokenized",
            file_format="loom",
        )
        return output_dir / "tokenized.dataset"

    def extract_baseline_embeddings(self, tokenized_path: Path, output_dir: Path) -> np.ndarray:
        """Extract baseline cell embeddings (before perturbation)."""
        embs = self.extractor.extract_embs(
            model_directory=self.model_variant,
            input_data_file=str(tokenized_path),
            output_directory=str(output_dir),
            output_prefix="baseline_embs",
        )
        return np.asarray(embs)

def in_silico_knockout(
    embedder: GeneformerEmbedder,
    tokenized_path: Path,
    target_gene_ensembl: str,
    output_dir: Path,
    max_ncells: int = 100,
) -> np.ndarray:
    """Perform in silico knockout for a specific gene and obtain perturbed cell embeddings.
    
    Uses the official Geneformer InSilicoPerturber API [3].
    """
    from geneformer import InSilicoPerturber
    
    perturber = InSilicoPerturber(
        perturb_type="delete",             # delete: remove gene, overexpress: move up
        perturb_rank_shift=None,            # None for delete
        genes_to_perturb=[target_gene_ensembl],
        combos=0,                            # Single gene
        anchor_gene=None,
        model_type="Pretrained",
        num_classes=0,
        emb_mode="cell",                     # Observe changes in cell embeddings
        cell_emb_style="mean_pool",
        filter_data=None,
        cell_states_to_model=None,
        max_ncells=max_ncells,               # Cell subset (for computational savings)
        emb_layer=-1,
        forward_batch_size=8,
        nproc=4,
    )
    
    perturbed_embs = perturber.perturb(
        model_directory=embedder.model_variant,
        input_data_file=str(tokenized_path),
        output_directory=str(output_dir),
        output_prefix=f"perturbed_{target_gene_ensembl}",
    )
    return np.asarray(perturbed_embs)

def cosine_shift(baseline_embs: np.ndarray, perturbed_embs: np.ndarray) -> np.ndarray:
    """Calculate cosine shift for each cell. Returns a shape of (N_cells,)."""
    baseline_norm = baseline_embs / (np.linalg.norm(baseline_embs, axis=1, keepdims=True) + 1e-8)
    perturbed_norm = perturbed_embs / (np.linalg.norm(perturbed_embs, axis=1, keepdims=True) + 1e-8)
    cos_sim = np.sum(baseline_norm * perturbed_norm, axis=1)
    return 1.0 - cos_sim  # Distance


def euclidean_shift(baseline_embs: np.ndarray, perturbed_embs: np.ndarray) -> np.ndarray:
    """Euclidean distance (auxiliary metric)."""
    return np.linalg.norm(perturbed_embs - baseline_embs, axis=1)


def screening_by_shift(
    embedder: GeneformerEmbedder,
    tokenized_path: Path,
    baseline_embs: np.ndarray,
    target_genes: list[str],       # List of Ensembl IDs
    output_dir: Path,
    max_ncells_per_gene: int = 100,
) -> dict[str, dict]:
    """Sequentially perform in silico KO for multiple genes and calculate shift statistics."""
    results = {}
    for i, gene in enumerate(target_genes):
        if i % 20 == 0:
            print(f"[{i}/{len(target_genes)}] {gene}")
        try:
            perturbed = in_silico_knockout(
                embedder, tokenized_path, gene, output_dir,
                max_ncells=max_ncells_per_gene,
            )
            cos = cosine_shift(baseline_embs[:len(perturbed)], perturbed)
            euc = euclidean_shift(baseline_embs[:len(perturbed)], perturbed)
            results[gene] = {
                "mean_cosine_shift": float(np.mean(cos)),
                "median_cosine_shift": float(np.median(cos)),
                "std_cosine_shift": float(np.std(cos)),
                "mean_euclidean_shift": float(np.mean(euc)),
                "n_cells": int(len(perturbed)),
            }
        except Exception as e:
            results[gene] = {"error": str(e)}
    
    return dict(sorted(
        results.items(),
        key=lambda x: -(x[1].get("mean_cosine_shift", 0.0) if "error" not in x[1] else 0.0),
    ))

import umap
import matplotlib.pyplot as plt


def visualize_perturbation_trajectory(
    baseline_embs: np.ndarray,
    perturbed_embs: np.ndarray,
    gene_symbol: str,
    output_path: str = "perturbation_umap.png",
    n_arrows: int = 20,
) -> None:
    """Visualize the cell state shift after knockout using UMAP."""
    combined = np.vstack([baseline_embs, perturbed_embs])
    reducer = umap.UMAP(n_components=2, n_neighbors=15, min_dist=0.1, random_state=42)
    embedding_2d = reducer.fit_transform(combined)
    
    n_cells = len(baseline_embs)
    baseline_2d = embedding_2d[:n_cells]
    perturbed_2d = embedding_2d[n_cells:]
    
    fig, ax = plt.subplots(figsize=(12, 8))
    ax.scatter(baseline_2d[:, 0], baseline_2d[:, 1],
               color="lightgray", alpha=0.5, s=20, label="baseline (WT)")
    ax.scatter(perturbed_2d[:, 0], perturbed_2d[:, 1],
               color="crimson", alpha=0.7, s=30, label=f"KO({gene_symbol})")
    
    # Arrows indicating the movement of sample cells
    n_arrows_actual = min(n_arrows, n_cells)
    idx = np.random.choice(n_cells, n_arrows_actual, replace=False)
    for i in idx:
        ax.annotate("", xy=perturbed_2d[i], xytext=baseline_2d[i],
                    arrowprops=dict(arrowstyle="->", color="black", alpha=0.4, lw=0.8))
    
    ax.set_title(f"Cell State Shift: In silico knockout of {gene_symbol}\n(n_cells={n_cells}, arrows={n_arrows_actual})")
    ax.set_xlabel("UMAP1")
    ax.set_ylabel("UMAP2")
    ax.legend()
    plt.tight_layout()
    plt.savefig(output_path, dpi=150)
    plt.close()

import pandas as pd
from scipy.stats import spearmanr


def load_replogle_ground_truth(
    replogle_h5ad_path: Path,
    baseline_ctrl_label: str = "non-targeting",
) -> pd.DataFrame:
    """Load the gene-wise KO effect from Replogle 2022 Perturb-seq.
    
    Download the K562 essential library from GEO GSE168191 and process it.
    Quantify how much the cell state of each target gene differs from the non-targeting control.
    """
    adata = sc.read_h5ad(replogle_h5ad_path)
    # Requires the 'target_gene' column
    if "target_gene" not in adata.obs.columns:
        raise ValueError("adata.obs['target_gene'] is required")
    
    # Control cell embeddings (e.g., PCA)
    if "X_pca" not in adata.obsm:
        sc.pp.pca(adata, n_comps=50)
    
    ctrl_mask = adata.obs["target_gene"] == baseline_ctrl_label
    ctrl_centroid = adata.obsm["X_pca"][ctrl_mask].mean(axis=0)
    
    # Distance between the centroid of each target gene and the control centroid
    effects = []
    for target in adata.obs["target_gene"].unique():
        if target == baseline_ctrl_label:
            continue
        target_mask = adata.obs["target_gene"] == target
        target_centroid = adata.obsm["X_pca"][target_mask].mean(axis=0)
        # Euclidean distance and cosine distance between centroids
        eucl_dist = float(np.linalg.norm(target_centroid - ctrl_centroid))
        cos_sim = float(np.dot(target_centroid, ctrl_centroid) / 
                        (np.linalg.norm(target_centroid) * np.linalg.norm(ctrl_centroid) + 1e-8))
        effects.append({
            "target_gene": target,
            "measured_shift": eucl_dist,
            "measured_cosine": 1.0 - cos_sim,
            "n_cells": int(target_mask.sum()),
        })
    return pd.DataFrame(effects)


def compare_with_wetlab(
    in_silico_results: dict[str, dict],
    wet_lab_df: pd.DataFrame,
    top_k: int = 20,
    gene_symbol_mapping: dict[str, str] | None = None,  # Ensembl to symbol mapping
) -> dict:
    """Compare in silico predictions with wet-lab measurements."""
    is_records = []
    for ensembl, stats in in_silico_results.items():
        if "error" in stats:
            continue
        symbol = gene_symbol_mapping.get(ensembl, ensembl) if gene_symbol_mapping else ensembl
        is_records.append({
            "target_gene": symbol,
            "silico_shift": stats["mean_cosine_shift"],
        })
    is_df = pd.DataFrame(is_records)
    merged = is_df.merge(wet_lab_df, on="target_gene", how="inner")
    print(f"Common genes: {len(merged)}")
    
    if len(merged) < 3:
        return {"error": "Insufficient common genes"}
    
    # Spearman correlation (shift magnitude)
    rho, pval = spearmanr(merged["silico_shift"], merged["measured_cosine"])
    
    # Recall at top-K
    is_topk = set(merged.nlargest(top_k, "silico_shift")["target_gene"])
    wl_topk = set(merged.nlargest(top_k, "measured_cosine")["target_gene"])
    recall_at_k = len(is_topk & wl_topk) / min(top_k, len(merged))
    
    # Precision at top-K
    precision_at_k = len(is_topk & wl_topk) / len(is_topk)
    
    return {
        "n_common_genes": len(merged),
        "spearman_rho": float(rho),
        "spearman_p": float(pval),
        f"recall_at_top{top_k}": float(recall_at_k),
        f"precision_at_top{top_k}": float(precision_at_k),
        "is_topk_genes": list(is_topk),
        "wl_topk_genes": list(wl_topk),
        "overlap_genes": list(is_topk & wl_topk),
    }

K562_ESSENTIAL_GENES_EXAMPLE = [
    # Example. In practice, select from DepMap, MAGeCK, or essential gene databases.
    "ENSG00000141510",  # TP53
    "ENSG00000012048",  # BRCA1
    "ENSG00000139618",  # BRCA2
    "ENSG00000186092",  # MYC (example)
    # ... 200 genes
]


def full_perturbation_pipeline(
    k562_h5ad_path: Path,
    essential_gene_ensembls: list[str],
    replogle_ref_path: Path | None,
    output_dir: Path,
    device: str = "cuda",
) -> dict:
    """Perform in silico screening of K562 essential genes and validate with Perturb-seq."""
    output_dir.mkdir(parents=True, exist_ok=True)
    
    print("[1/6] K562 Data QC")
    adata = load_and_qc(k562_h5ad_path)
    
    print("[2/6] Load Geneformer")
    embedder = GeneformerEmbedder(device=device)
    
    print("[3/6] Tokenize & Baseline Embedding")
    tokenized = embedder.tokenize_adata(adata, output_dir / "tokens")
    baseline = embedder.extract_baseline_embeddings(tokenized, output_dir / "baseline")
    print(f"      baseline shape={baseline.shape}")
    
    print(f"[4/6] In silico KO of {len(essential_gene_ensembls)} genes")
    silico_results = screening_by_shift(
        embedder, tokenized, baseline, essential_gene_ensembls,
        output_dir / "perturbations", max_ncells_per_gene=50,
    )
    import json
    with open(output_dir / "silico_shifts.json", "w") as f:
        json.dump(silico_results, f, indent=2, ensure_ascii=False)
    
    print("[5/6] Visualize Top Gene UMAP Trajectories")
    top_genes = [g for g in list(silico_results.keys())[:5] if "error" not in silico_results[g]]
    for gene in top_genes:
        perturbed_dir = output_dir / "perturbations" / f"perturbed_{gene}"
        # ... reload perturbed results from embedder or cache
        pass  # (In practice, use caching to avoid recalculation for visualization)
    
    print("[6/6] Compare with Perturb-seq Measurements")
    validation = {}
    if replogle_ref_path:
        wet_lab_df = load_replogle_ground_truth(replogle_ref_path)
        validation = compare_with_wetlab(silico_results, wet_lab_df, top_k=20)
        with open(output_dir / "validation.json", "w") as f:
            json.dump(validation, f, indent=2, ensure_ascii=False)
        print(f"      Spearman ρ={validation.get('spearman_rho', 0):.3f}, "
              f"Recall@20={validation.get('recall_at_top20', 0):.2f}")
    
    return {"silico_ranking": silico_results, "validation": validation}

## Performance, Cost, and Known Failure Cases

### Performance Reference (Using Public Benchmarks)

| Model | Benchmark (Perturb-seq recall@20) | Spearman ρ | Source |
|------|:----------------------------:|:----------:|------|
| Random baseline | K562 essential | 0.10 | Baseline |
| GENIE3 (correlation-based) | K562 essential | 0.25 | Legacy Method |
| scGPT in silico KO | K562 essential | 0.45~0.55 | Cui et al., Nat Methods 2024 [2] |
| Geneformer in silico KO | Replogle 2022 subset | 0.50~0.65 | Theodoris et al., Nature 2023 [3] |
| scFoundation | Multi-tissue | 0.55~0.70 | Hao et al., Nat Methods 2024 [4] |
| UCE | Cross-species | 0.50~0.60 | Rosen et al. 2024 [5] |
| Ensemble (scGPT + Geneformer) | K562 | ~0.70 (Estimated) | Community Benchmark |

### Estimated Cost for Learners' Reproduction

- API cost: 0 (when using a local GPU).
- Learners without a powerful GPU should refer to the cloud on-demand hourly pricing (Geneformer 95M can run on 24GB VRAM).
- Screening 200 genes takes approximately 30-60 minutes.
- Data download: Geneformer weights 4GB + Replogle Perturb-seq subset 5~10GB.

### 5 Known Failure Cases (Collected from the Community and Papers)

1. **Discrepancy between in silico KO and wet-lab results (Specific cell types)**
   Symptoms: In silico KO predictions yield meaningless values for cell types not present in the pre-training data (e.g., a specific subtype of brain neurons).
   Cause: Limitation of the coverage of the pre-training data of the foundation model.
   Mitigation: (a) Fine-tune the model with data containing the target cell type, (b) Ensemble multiple foundation models (Geneformer + scGPT), (c) Filter predictions based on confidence (when the magnitude of the embedding shift is very small), (d) K562 is included in the pre-training data, making it relatively stable.
   Source: Kedzierska et al. "Assessing the limits of zero-shot foundation models in single-cell biology." bioRxiv 2023 [6].

2. **Gene Ensembl ID mapping errors**
   Symptoms: Failure to map gene symbols (TP53) to Ensembl IDs (ENSG00000141510) → in silico KO is applied to the wrong gene.
   Cause: Gene symbols have aliases, and Geneformer is trained on a specific Ensembl version (e.g., GRCh38.p13).
   Mitigation: (a) Use `pyensembl` or `mygene.info` to accurately map and specify the version, (b) Pin the Ensembl version, (c) Log and skip genes with mapping failures, (d) Check for updated Geneformer retraining versions.
   Source: Geneformer GitHub Issues [7].

3. **Batch effects contaminating in silico shift**
   Symptoms: Analyzing cells from different batches together results in batch effects being misinterpreted as shifts, leading to incorrect gene effect interpretations.
   Cause: Pre-QC and normalization do not completely remove batch effects.
   Mitigation: (a) Batch effect correction (Harmony, scVI, Scanorama), (b) Perform in silico KO only within a single batch, (c) Use the shift of control genes (neutral controls) as the baseline noise to calculate signal-to-noise ratio, (d) Verify the reproducibility of shift values across multiple batches.
   Source: Peidli S et al. "scPerturb: harmonized single-cell perturbation data." Nat Methods 2024 [8].

4. **Gap between rank drop protocol and actual KO**
   Symptoms: Rank drop simulations differ from real gRNA-Cas9 KO results.
   Cause: Rank drop simply shifts the expression ranking, while actual KO reflects protein, complex, downstream cascade, and temporal dynamics.
   Mitigation: (a) Interpret only genes with large shifts as strong signals, (b) Try multiple perturbation modes of Geneformer (delete, overexpress, knockdown), (c) Use the results only for hypothesis generation and require wet-lab validation, (d) Quantitatively assess consistency with the Perturb-seq benchmark.
   Source: Theodoris et al. Nature 2023 discussion [3]; Kedzierska 2023 [6].

5. **Burden of downloading and storing large datasets**
   Symptoms: The entire Replogle 2022 Perturb-seq dataset is tens of GB, and the K562 essential subset is also 5-10GB.
   Cause: Raw counts of single-cell data are large, even as a sparse matrix.
   Mitigation: (a) Use subsets (essential gene 200 subset only), (b) Download only the necessary samples from GEO, (c) Store with compression (h5ad + zstd), (d) Utilize academic cloud resources (e.g., CZI Chan Zuckerberg Initiative BioHub).
   Source: GEO GSE168191 Replogle 2022 [1].

## Expansion Ideas

- **Overexpression simulation:** Simulate overexpression by ranking up instead of ranking down (e.g., transcription factors, tumor suppressor genes).
- **Multi-gene combinations:** Simultaneously knock out two genes (screen for synthetic lethality candidates).
- **Drug response prediction:** Compare the cell shift of drug target KO with actual drug treatment data → drug repurposing (link to sections 07 and 08).
- **Cross-species transfer:** Apply patterns learned from mouse data to human data (using UCE).
- **Live scRNA-seq integration:** Stream real-time experimental data → compare in silico predictions with real-time data.
- **Detailed prediction of HIF network and signaling pathways:** KO combinations of specific signaling pathways.

## Next Section

- Section 04 `crispr-guide-scoring`: Design gRNAs for genes screened in silico (entry into K562 wet-lab).
- Section 07 `drug-target-gnn`: Utilize gene KO prediction to prioritize drug targets.
- Section 13 `protein-design-multimodal`: Design artificial activator proteins to replace KO genes.
- Section 14 `bio-mcp-agent`: Expose in silico perturbation as an MCP tool → "Extract genes with large shifts in K562" for autonomous execution.

## References

1. Replogle JM, Saunders RA, Pogson AN, et al. "Mapping information-rich genotype-phenotype landscapes with genome-scale Perturb-seq." Cell 2022. `https://www.cell.com/cell/fulltext/S0092-8674(22)00597-9` · GEO GSE168191
2. Cui H, Wang C, Maan H, et al. "scGPT: toward building a foundation model for single-cell multi-omics using generative AI." Nature Methods 2024. `https://www.nature.com/articles/s41592-024-02201-0`
3. Theodoris CV, Xiao L, Chopra A, et al. "Transfer learning enables predictions in network biology (Geneformer)." Nature 2023. `https://www.nature.com/articles/s41586-023-06139-9`
4. Hao M, Gong J, Zeng X, et al. "Large-scale foundation model on single-cell transcriptomics (scFoundation)." Nature Methods 2024. `https://www.nature.com/articles/s41592-024-02305-7`
5. Rosen Y, Roohani Y, Agarwal A, et al. "Universal Cell Embeddings: A Foundation Model for Cell Biology (UCE)." bioRxiv 2024. `https://www.biorxiv.org/content/10.1101/2023.11.28.568918v2`
6. Kedzierska KZ, Crawford L, Amini AP, Lu AX. "Assessing the limits of zero-shot foundation models in single-cell biology." bioRxiv 2023. `https://www.biorxiv.org/content/10.1101/2023.10.16.561085`
7. Geneformer GitHub Issues: `https://huggingface.co/ctheodoris/Geneformer/discussions`
8. Peidli S, Green TD, Shen C, et al. "scPerturb: harmonized single-cell perturbation data." Nature Methods 2024. `https://www.nature.com/articles/s41592-023-02144-y`
9. Geneformer HuggingFace: `https://huggingface.co/ctheodoris/Geneformer`
10. scGPT GitHub: `https://github.com/bowang-lab/scGPT`
11. CELLxGENE (CZI): `https://cellxgene.cziscience.com/`
12. scanpy: `https://scanpy.readthedocs.io/`
13. anndata: `https://anndata.readthedocs.io/`
14. Human Cell Atlas: `https://www.humancellatlas.org/`
15. UMAP-learn: `https://umap-learn.readthedocs.io/`
16. helical SDK (Geneformer/scGPT/UCE 통합): `https://github.com/helicalAI/helical`
17. DepMap (essential gene DB): `https://depmap.org/`
18. MAGeCK (CRISPR screen 분석): `https://sourceforge.net/p/mageck/wiki/Home/`
19. Harmony batch correction: `https://github.com/immunogenomics/harmony`
20. scVI (batch correction + generative model): `https://scvi-tools.org/`