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:
    """K562 데이터셋 표준 QC."""
    adata = sc.read_h5ad(h5ad_path)
    print(f"로딩: cells {adata.n_obs}, genes {adata.n_vars}")
    
    # 미토콘드리아 · 리보솜 유전자 태그
    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)
    
    # 필터링
    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()
    
    # 정규화 · log 변환
    sc.pp.normalize_total(adata, target_sum=1e4)
    sc.pp.log1p(adata)
    
    # HVG · scale (Geneformer는 raw counts 사용하므로 별도 layer 유지)
    if "raw_counts" not in adata.layers:
        adata.layers["raw_counts"] = adata.X.copy()  # Geneformer 입력용
    
    print(f"QC 후: cells {adata.n_obs}, genes {adata.n_vars}")
    return adata

import torch


class GeneformerEmbedder:
    """Geneformer 95M 사전학습 모델 wrapping.
    
    Geneformer 공식 저장소의 EmbExtractor · TranscriptomeTokenizer 사용.
    """

    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는 자체 tokenizer + BertForMaskedLM 구조
        # 실전 코드는 helical 또는 공식 저장소 예시 참조
        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",       # 세포 임베딩 (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:
        """AnnData → Geneformer 토큰 파일 (.dataset).
        
        Ensembl gene ID (adata.var["ensembl_id"]) 필수.
        """
        if "ensembl_id" not in adata.var.columns:
            # 심볼 → Ensembl 매핑 (mygene · pyensembl)
            raise ValueError("adata.var['ensembl_id'] 필수. mygene 또는 pyensembl로 사전 매핑.")
        
        output_dir.mkdir(parents=True, exist_ok=True)
        # anndata → loom 변환 후 tokenize (Geneformer 공식 워크플로)
        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:
        """기준 세포 임베딩 (섭동 전)."""
        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:
    """특정 유전자 in silico knockout → 섭동 세포 임베딩.
    
    Geneformer 공식 InSilicoPerturber API [3].
    """
    from geneformer import InSilicoPerturber
    
    perturber = InSilicoPerturber(
        perturb_type="delete",             # delete: 유전자 제거, overexpress: 상위 이동
        perturb_rank_shift=None,            # delete는 None
        genes_to_perturb=[target_gene_ensembl],
        combos=0,                            # 단일 유전자
        anchor_gene=None,
        model_type="Pretrained",
        num_classes=0,
        emb_mode="cell",                     # 세포 임베딩 변화 관찰
        cell_emb_style="mean_pool",
        filter_data=None,
        cell_states_to_model=None,
        max_ncells=max_ncells,               # 세포 subset (계산 절약)
        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:
    """세포별 코사인 변위. shape: (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  # 거리


def euclidean_shift(baseline_embs: np.ndarray, perturbed_embs: np.ndarray) -> np.ndarray:
    """유클리드 거리 (보조 지표)."""
    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],       # Ensembl ID 리스트
    output_dir: Path,
    max_ncells_per_gene: int = 100,
) -> dict[str, dict]:
    """여러 유전자 순차 in silico KO → shift 통계."""
    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:
    """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})")
    
    # 샘플 세포의 이동 화살표
    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"세포 상태 이동: {gene_symbol} in silico knockout\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:
    """Replogle 2022 Perturb-seq에서 유전자별 KO effect 로딩.
    
    GEO GSE168191에서 K562 essential library 다운로드 후 처리.
    각 target_gene의 세포 상태가 non-targeting control 대비 얼마나 이동했는지 정량.
    """
    adata = sc.read_h5ad(replogle_h5ad_path)
    # target_gene column 필요
    if "target_gene" not in adata.obs.columns:
        raise ValueError("adata.obs['target_gene'] 필수")
    
    # Control 세포 임베딩 (예: 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)
    
    # 각 target 세포의 centroid와 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)
        # centroid 간 유클리드 거리 · cosine 거리
        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 → symbol
) -> dict:
    """in silico 예측과 wet-lab 실측 비교."""
    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"공통 유전자: {len(merged)}")
    
    if len(merged) < 3:
        return {"error": "공통 유전자 부족"}
    
    # Spearman 상관 (shift 크기)
    rho, pval = spearmanr(merged["silico_shift"], merged["measured_cosine"])
    
    # 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 = [
    # 예시. 실제로는 DepMap · MAGeCK · essential gene DB에서 선별.
    "ENSG00000141510",  # TP53
    "ENSG00000012048",  # BRCA1
    "ENSG00000139618",  # BRCA2
    "ENSG00000186092",  # MYC (예시)
    # ... 200개
]


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:
    """K562 essential gene in silico 스크리닝 · Perturb-seq 검증."""
    output_dir.mkdir(parents=True, exist_ok=True)
    
    print("[1/6] K562 데이터 QC")
    adata = load_and_qc(k562_h5ad_path)
    
    print("[2/6] Geneformer 로딩")
    embedder = GeneformerEmbedder(device=device)
    
    print("[3/6] 토큰화 · baseline 임베딩")
    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] {len(essential_gene_ensembls)}개 유전자 in silico KO 스크리닝")
    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] 상위 유전자 UMAP 궤적 시각화")
    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}"
        # ... embedder에서 perturbed 결과 재로드 or 캐싱
        pass  # (실전에서는 caching 활용해 재계산 없이 시각화)
    
    print("[6/6] Perturb-seq 실측 비교")
    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}