#!/usr/bin/env python
"""Comprehensive wrap-up analysis: address weaknesses, add rigor.

No retraining — uses existing fitted results + re-analyzes data.

1. Fair comparison: analytical vs DeepPTR on SAME 300 genes
2. Bootstrap CIs on half-life correlations
3. Validate PT-specific genes against eCLIP RBP targets
4. Examine sci-fate tautology honestly
5. Sparsity analysis: gamma quality vs unspliced detection rate
6. CI coverage breakdown: where does the posterior fail?
7. ARE/NMD enrichment of PT-specific genes
8. Honest limitations table

All results saved to output/wrapup/.
"""

from __future__ import annotations

import os
os.environ["OMP_NUM_THREADS"] = "4"
os.environ["MKL_NUM_THREADS"] = "4"
os.environ["OPENBLAS_NUM_THREADS"] = "4"

import json
import sys
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from scipy import stats

import torch
torch.set_num_threads(4)

sys.path.insert(0, str(Path(__file__).parent))
from _common import set_figure_style

import scptr

OUTPUT_DIR = Path(__file__).parent.parent / "output" / "wrapup"
DATA_DIR = Path(scptr.benchmark.__file__).parent / "data"


def save_fig(fig, name, subdir="figures"):
    if fig is None:
        return
    out_dir = OUTPUT_DIR / subdir
    out_dir.mkdir(parents=True, exist_ok=True)
    path = out_dir / f"{name}.png"
    fig.savefig(path, dpi=150, bbox_inches="tight")
    plt.close(fig)
    print(f"  Saved: {path}")


def ensure_dirs():
    for sub in ("figures", "results"):
        (OUTPUT_DIR / sub).mkdir(parents=True, exist_ok=True)


def select_top_genes(adata, n_top=300):
    from scipy.sparse import issparse
    u = adata.layers["unspliced"]
    if issparse(u):
        u = np.asarray(u.todense())
    u = np.asarray(u, dtype=np.float32)
    score = u.sum(axis=0) * (u > 0).mean(axis=0)
    top_idx = np.sort(np.argsort(score)[::-1][:n_top])
    return adata.var_names[top_idx].tolist()


def prepare_analytical(adata_loader):
    adata = adata_loader()
    scptr.pp.filter_genes(adata)
    scptr.pp.normalize_layers(adata)
    scptr.pp.neighbors(adata, n_neighbors=30)
    scptr.pp.smooth_layers(adata)
    scptr.tl.estimate_beta(adata)
    scptr.tl.estimate_gamma(adata)
    return adata


# ============================================================================
# 1. FAIR COMPARISON: SAME 300 GENES
# ============================================================================

def analysis_fair_comparison(adata_an, dataset_name, top_genes):
    """Compare half-life correlation using analytical gamma on the SAME 300 genes."""
    print(f"\n{'=' * 60}")
    print(f"1. FAIR COMPARISON: Same 300 genes ({dataset_name})")
    print("=" * 60)

    # Load previous DeepPTR results
    prev_file = Path(__file__).parent.parent / "output" / "deep_benchmark" / "results" / f"{dataset_name}_benchmark.json"
    if prev_file.exists():
        with open(prev_file) as f:
            prev = json.load(f)
    else:
        prev = {}

    # Analytical on ALL genes
    hl_mouse = scptr.datasets.herzog2017_halflives()
    hl_human = scptr.datasets.schofield2018_halflives()

    gamma_all = np.median(adata_an.layers["gamma"], axis=0)

    # Analytical on SAME 300 genes
    gene_mask = np.isin(adata_an.var_names, top_genes)
    gamma_300 = gamma_all.copy()
    gamma_300[~gene_mask] = 0  # zero out genes not in top-300

    results = {}
    for ref_name, hl_df in [("mouse", hl_mouse), ("human", hl_human)]:
        # Full analytical
        corr_full = scptr.benchmark.correlate_with_halflives(adata_an, hl_df)

        # Analytical restricted to 300 genes (create temp adata)
        adata_300 = adata_an[:, top_genes].copy()
        # Need gamma layer
        an_300_idx = [list(adata_an.var_names).index(g) for g in top_genes if g in adata_an.var_names]
        adata_300.layers["gamma"] = adata_an.layers["gamma"][:, an_300_idx]
        corr_300 = scptr.benchmark.correlate_with_halflives(adata_300, hl_df)

        # DeepPTR from previous results
        hl_key = "mouse_herzog" if ref_name == "mouse" else "human_schofield"
        dp_r = prev.get("halflife", {}).get(hl_key, {}).get("deepptr", {}).get("spearman_r", np.nan)
        dp_n = prev.get("halflife", {}).get(hl_key, {}).get("deepptr", {}).get("n_genes", 0)

        results[ref_name] = {
            "analytical_all": {"r": corr_full["spearman_r"], "n": corr_full["n_genes"]},
            "analytical_300": {"r": corr_300["spearman_r"], "n": corr_300["n_genes"]},
            "deepptr_300": {"r": dp_r, "n": dp_n},
        }

        print(f"\n  {ref_name}:")
        print(f"    Analytical (all {adata_an.n_vars} genes): r={corr_full['spearman_r']:.4f} (n={corr_full['n_genes']})")
        print(f"    Analytical (same 300 genes):  r={corr_300['spearman_r']:.4f} (n={corr_300['n_genes']})")
        print(f"    DeepPTR    (same 300 genes):  r={dp_r:.4f} (n={dp_n})")

    return results


# ============================================================================
# 2. BOOTSTRAP CONFIDENCE INTERVALS
# ============================================================================

def analysis_bootstrap_ci(adata_an, dataset_name, n_boot=1000):
    """Bootstrap CIs on half-life correlations."""
    print(f"\n{'=' * 60}")
    print(f"2. BOOTSTRAP CIs ({dataset_name})")
    print("=" * 60)

    hl_human = scptr.datasets.schofield2018_halflives()
    hl_s = hl_human.set_index("gene_symbol")["half_life_hours"]

    gamma_med = np.median(adata_an.layers["gamma"], axis=0)
    gamma_s = pd.Series(gamma_med, index=adata_an.var_names)

    # Case-insensitive match
    gamma_upper = {g.upper(): g for g in gamma_s.index}
    hl_upper = {g.upper(): g for g in hl_s.index if isinstance(g, str)}
    shared_upper = set(gamma_upper.keys()) & set(hl_upper.keys())

    g_vals = np.array([gamma_s[gamma_upper[u]] for u in shared_upper], dtype=float)
    h_vals = np.array([hl_s[hl_upper[u]] for u in shared_upper], dtype=float)

    valid = np.isfinite(g_vals) & np.isfinite(h_vals) & (g_vals > 0) & (h_vals > 0)
    g_vals, h_vals = g_vals[valid], h_vals[valid]
    n = len(g_vals)

    # Point estimate
    sp_r, _ = stats.spearmanr(g_vals, h_vals)

    # Bootstrap
    rng = np.random.RandomState(42)
    boot_rs = np.zeros(n_boot)
    for i in range(n_boot):
        idx = rng.choice(n, size=n, replace=True)
        boot_rs[i], _ = stats.spearmanr(g_vals[idx], h_vals[idx])

    ci_lo, ci_hi = np.percentile(boot_rs, [2.5, 97.5])
    se = np.std(boot_rs)

    print(f"  Spearman r = {sp_r:.4f} (n={n})")
    print(f"  95% CI: [{ci_lo:.4f}, {ci_hi:.4f}]")
    print(f"  Bootstrap SE: {se:.4f}")

    result = {
        "spearman_r": float(sp_r),
        "n_genes": n,
        "ci_95_lo": float(ci_lo),
        "ci_95_hi": float(ci_hi),
        "bootstrap_se": float(se),
    }
    return result


# ============================================================================
# 3. eCLIP VALIDATION OF PT-SPECIFIC GENES
# ============================================================================

def analysis_eclip_validation(dataset_name):
    """Check if PT-specific genes are enriched for eCLIP RBP targets."""
    print(f"\n{'=' * 60}")
    print(f"3. eCLIP VALIDATION ({dataset_name})")
    print("=" * 60)

    # Load PT-specific genes from previous analysis
    adv_file = Path(__file__).parent.parent / "output" / "deep_advantages" / "results" / f"{dataset_name}_advantages.json"
    if not adv_file.exists():
        print("  [SKIP] No advantage results found")
        return None

    with open(adv_file) as f:
        adv = json.load(f)

    pt_genes = adv.get("disentanglement", {}).get("pt_specific_genes", [])
    if not pt_genes:
        print("  [SKIP] No PT-specific genes")
        return None

    # Load eCLIP targets
    eclip = pd.read_csv(DATA_DIR / "eclip_targets.csv")
    eclip_targets = set(eclip["target_gene"].str.upper())
    eclip_by_rbp = eclip.groupby("rbp")["target_gene"].apply(lambda x: set(x.str.upper())).to_dict()

    # Test: are PT-specific genes enriched for eCLIP targets?
    pt_upper = set(g.upper() for g in pt_genes)

    # Also load the full gene list for background
    # Use all 300 DeepPTR genes as background
    pt_de_genes = [g["gene"] for g in adv.get("disentanglement", {}).get("top_pt_de_genes", [])]
    all_genes_upper = pt_upper | set(g.upper() for g in pt_de_genes)

    # If we don't have enough background, we can't do enrichment
    # Let's just count overlap
    pt_in_eclip = pt_upper & eclip_targets
    frac_pt = len(pt_in_eclip) / max(len(pt_upper), 1)

    print(f"  PT-specific genes: {len(pt_genes)}")
    print(f"  In eCLIP database: {len(pt_in_eclip)} ({frac_pt*100:.0f}%)")
    if pt_in_eclip:
        print(f"  Validated genes: {sorted(pt_in_eclip)[:20]}")

    # Per-RBP enrichment: which RBPs target PT-specific genes?
    rbp_hits = {}
    for rbp, targets in eclip_by_rbp.items():
        overlap = pt_upper & targets
        if overlap:
            rbp_hits[rbp] = sorted(overlap)

    if rbp_hits:
        print(f"\n  RBPs targeting PT-specific genes:")
        for rbp in sorted(rbp_hits, key=lambda x: len(rbp_hits[x]), reverse=True)[:10]:
            print(f"    {rbp}: {len(rbp_hits[rbp])} targets — {rbp_hits[rbp][:5]}")

    # Fisher's exact test: are PT genes more likely to be eCLIP targets than random?
    # Background: all genes in the dataset
    result = {
        "n_pt_genes": len(pt_genes),
        "n_in_eclip": len(pt_in_eclip),
        "frac_in_eclip": frac_pt,
        "validated_genes": sorted(pt_in_eclip),
        "rbp_hits": {k: v for k, v in sorted(rbp_hits.items(), key=lambda x: len(x[1]), reverse=True)[:15]},
    }

    return result


# ============================================================================
# 4. SCI-FATE TAUTOLOGY ANALYSIS
# ============================================================================

def analysis_scifate_tautology():
    """Honestly examine the sci-fate tautology concern.

    gamma ∝ beta * Mu / Ms ∝ new / old (approximately)
    ground truth = new / old

    How much of the r=0.99 is structural vs learned?
    """
    print(f"\n{'=' * 60}")
    print("4. SCI-FATE TAUTOLOGY ANALYSIS")
    print("=" * 60)

    import gzip
    from scipy.io import mmread
    from scipy.sparse import csc_matrix

    CACHE_DIR = Path.home() / ".cache" / "scptr" / "scifate"
    if not CACHE_DIR.exists():
        print("  [SKIP] sci-fate data not cached")
        return None

    # Load data
    cell_ann = pd.read_csv(CACHE_DIR / "GSM3770930_A549_cell_annotate.txt.gz", compression="gzip")
    gene_ann = pd.read_csv(CACHE_DIR / "GSM3770930_A549_gene_annotate.txt.gz", compression="gzip")

    with gzip.open(CACHE_DIR / "GSM3770930_A549_gene_count.txt.gz", "rb") as f:
        total_mat = csc_matrix(mmread(f)).T
    with gzip.open(CACHE_DIR / "GSM3770930_A549_gene_count_newly_synthesised.txt.gz", "rb") as f:
        new_mat = csc_matrix(mmread(f)).T

    total = np.asarray(total_mat.todense())
    new = np.asarray(new_mat.todense())
    old = total - new

    mean_new = new.mean(axis=0)
    mean_old = old.mean(axis=0)
    mean_total = total.mean(axis=0)

    reliable = (mean_total >= 0.5) & (mean_old > 0.1)
    gt_ratio = np.full(total.shape[1], np.nan)
    gt_ratio[reliable] = mean_new[reliable] / mean_old[reliable]

    # The mapping: unspliced=new, spliced=old
    # So gamma = beta * mean(new) / mean(old) [approximately, after smoothing]
    # And ground truth = mean(new) / mean(old)
    # Therefore gamma ≈ beta * ground_truth
    # Correlation(gamma, ground_truth) ≈ Correlation(beta * GT, GT) = high if beta has low variance

    # Compute the "trivial baseline": raw ratio new/old (no model needed)
    trivial_ratio = np.full(total.shape[1], np.nan)
    trivial_ratio[reliable] = mean_new[reliable] / mean_old[reliable]

    # Now run the pipeline to get actual gamma
    import anndata as ad
    keep = mean_total >= 0.5
    if "gene_type" in gene_ann.columns:
        is_pc = gene_ann["gene_type"] == "protein_coding"
        keep = keep & is_pc.values

    gene_ann_indexed = gene_ann.set_index("gene_id")
    adata = ad.AnnData(
        X=total[:, keep].astype(np.float32),
        obs=cell_ann.set_index("sample"),
        var=gene_ann_indexed.iloc[keep].copy(),
    )
    adata.layers["unspliced"] = new[:, keep].astype(np.float32)
    adata.layers["spliced"] = old[:, keep].astype(np.float32)
    adata.var_names = adata.var["gene_short_name"].values
    adata.var_names_make_unique()

    scptr.pp.filter_genes(adata, min_unspliced_counts=1, min_unspliced_cells=1)
    scptr.pp.normalize_layers(adata)
    scptr.pp.neighbors(adata, n_neighbors=30)
    scptr.pp.smooth_layers(adata)
    scptr.tl.estimate_beta(adata)
    scptr.tl.estimate_gamma(adata)

    gamma_med = np.median(adata.layers["gamma"], axis=0)
    beta_vals = adata.var["beta"].values

    # Match with ground truth using case-insensitive matching
    gamma_s = pd.Series(gamma_med, index=adata.var_names)
    beta_s = pd.Series(beta_vals, index=adata.var_names)

    # Build ground truth series indexed by gene short names (deduplicated)
    gene_names_raw = gene_ann["gene_short_name"].values
    gt_dict = {}
    for i, gn in enumerate(gene_names_raw):
        if isinstance(gn, str) and reliable[i] and gn not in gt_dict:
            gt_dict[gn] = gt_ratio[i]
    gt_s = pd.Series(gt_dict)

    shared = gamma_s.index.intersection(gt_s.dropna().index)
    g = gamma_s[shared].values.astype(float)
    t = gt_s[shared].values.astype(float)
    b = beta_s[shared].values.astype(float)

    valid = np.isfinite(g) & np.isfinite(t) & (g > 0) & (t > 0) & np.isfinite(b)
    g, t, b = g[valid], t[valid], b[valid]

    # Correlations
    r_gamma_gt, _ = stats.spearmanr(g, t)  # gamma vs ground truth
    r_trivial, _ = stats.spearmanr(t, t)  # trivial = 1.0

    # Partial out beta: correlation of gamma with GT controlling for beta
    # gamma ≈ beta * GT, so gamma/beta ≈ GT
    gamma_over_beta = g / (b + 1e-8)
    r_residual, _ = stats.spearmanr(gamma_over_beta, t)

    # How much does beta vary?
    beta_cv = np.std(b) / np.mean(b)

    # Correlation of beta with gamma (if beta is constant, gamma ∝ GT exactly)
    r_beta_gamma, _ = stats.spearmanr(b, g)

    print(f"  n genes: {len(g)}")
    print(f"  gamma vs ground truth:   r = {r_gamma_gt:.4f}")
    print(f"  gamma/beta vs GT:        r = {r_residual:.4f}")
    print(f"  beta CV:                 {beta_cv:.4f}")
    print(f"  beta vs gamma:           r = {r_beta_gamma:.4f}")
    print(f"\n  Interpretation:")
    print(f"    gamma = beta * (Mu/Ms) ≈ beta * (new/old) = beta * GT")
    print(f"    Since beta CV = {beta_cv:.2f}, beta adds {'modest' if beta_cv < 0.5 else 'substantial'} variation")
    print(f"    After dividing out beta, residual r = {r_residual:.4f}")
    print(f"    → The r={r_gamma_gt:.3f} correlation is {'largely' if r_residual > 0.95 else 'partially'} "
          f"tautological")

    # What scPTR ADDS beyond the trivial ratio: the smoothing, beta correction,
    # and clipping — test if these improve the correlation
    # Raw ratio (no smoothing, no beta): just new/old per cell, median across cells
    raw_ratio = np.median(new[:, keep], axis=0) / np.clip(np.median(old[:, keep], axis=0), 1e-8, None)
    raw_s = pd.Series(raw_ratio, index=adata.var_names[:len(raw_ratio)])
    shared2 = raw_s.index.intersection(gt_s.dropna().index)
    r_raw_vals = raw_s[shared2].values.astype(float)
    t_raw_vals = gt_s[shared2].values.astype(float)
    v2 = np.isfinite(r_raw_vals) & np.isfinite(t_raw_vals) & (r_raw_vals > 0) & (t_raw_vals > 0)
    if v2.sum() > 3:
        r_raw, _ = stats.spearmanr(r_raw_vals[v2], t_raw_vals[v2])
        print(f"\n  Raw median(new)/median(old) vs GT: r = {r_raw:.4f} (n={v2.sum()})")
        print(f"  scPTR pipeline adds:               Δr = {r_gamma_gt - r_raw:.4f}")
    else:
        r_raw = np.nan

    result = {
        "r_gamma_gt": float(r_gamma_gt),
        "r_gamma_over_beta_gt": float(r_residual),
        "r_raw_ratio_gt": float(r_raw) if not np.isnan(r_raw) else None,
        "beta_cv": float(beta_cv),
        "r_beta_gamma": float(r_beta_gamma),
        "n_genes": len(g),
        "tautology_severity": "high" if r_residual > 0.98 else "moderate" if r_residual > 0.90 else "low",
    }

    return result


# ============================================================================
# 5. SPARSITY ANALYSIS
# ============================================================================

def analysis_sparsity(adata_an, dataset_name):
    """Does gamma quality depend on unspliced detection rate?"""
    print(f"\n{'=' * 60}")
    print(f"5. SPARSITY ANALYSIS ({dataset_name})")
    print("=" * 60)

    from scipy.sparse import issparse

    u = adata_an.layers["unspliced"]
    if issparse(u):
        u = np.asarray(u.todense())
    u = np.asarray(u)

    # Per-gene: fraction of cells with unspliced > 0
    frac_detected = (u > 0).mean(axis=0)

    gamma_med = np.median(adata_an.layers["gamma"], axis=0)

    # Half-life correlation stratified by detection rate
    hl_human = scptr.datasets.schofield2018_halflives()
    hl_s = hl_human.set_index("gene_symbol")["half_life_hours"]

    gamma_upper = {g.upper(): i for i, g in enumerate(adata_an.var_names)}
    hl_upper = {g.upper(): g for g in hl_s.index if isinstance(g, str)}
    shared = set(gamma_upper.keys()) & set(hl_upper.keys())

    g_idx = np.array([gamma_upper[u] for u in shared])
    h_vals = np.array([hl_s[hl_upper[u]] for u in shared], dtype=float)
    g_vals = gamma_med[g_idx]
    det_vals = frac_detected[g_idx]

    valid = np.isfinite(g_vals) & np.isfinite(h_vals) & (g_vals > 0) & (h_vals > 0)
    g_vals, h_vals, det_vals = g_vals[valid], h_vals[valid], det_vals[valid]

    # Stratify by detection quartile
    quartiles = np.percentile(det_vals, [25, 50, 75])
    bins = [
        ("Q1 (lowest)", det_vals <= quartiles[0]),
        ("Q2", (det_vals > quartiles[0]) & (det_vals <= quartiles[1])),
        ("Q3", (det_vals > quartiles[1]) & (det_vals <= quartiles[2])),
        ("Q4 (highest)", det_vals > quartiles[2]),
    ]

    records = []
    print(f"\n  Half-life correlation by unspliced detection rate:")
    for label, mask in bins:
        if mask.sum() < 10:
            continue
        sp_r, _ = stats.spearmanr(g_vals[mask], h_vals[mask])
        records.append({
            "quartile": label,
            "n_genes": int(mask.sum()),
            "spearman_r": float(sp_r),
            "median_detection": float(np.median(det_vals[mask])),
        })
        print(f"    {label}: r={sp_r:.4f} (n={mask.sum()}, median det={np.median(det_vals[mask]):.2f})")

    # Overall correlation: detection rate vs |gamma - halflife rank correlation|
    r_det, p_det = stats.spearmanr(det_vals, np.abs(g_vals))
    print(f"\n  Detection rate vs |gamma|: r={r_det:.4f} (p={p_det:.2e})")

    # Plot
    fig, axes = plt.subplots(1, 2, figsize=(12, 5))

    ax = axes[0]
    for rec in records:
        ax.bar(rec["quartile"], abs(rec["spearman_r"]), color="steelblue", alpha=0.7)
    ax.set_ylabel("|Spearman r| with half-life")
    ax.set_title(f"{dataset_name}: Half-life r by detection rate")
    ax.set_xticklabels([r["quartile"] for r in records], rotation=30, ha="right")

    ax = axes[1]
    ax.scatter(det_vals, g_vals, alpha=0.1, s=3, c="steelblue")
    ax.set_xlabel("Unspliced detection rate")
    ax.set_ylabel("Median gamma")
    ax.set_title(f"Detection rate vs gamma (r={r_det:.3f})")

    fig.tight_layout()
    save_fig(fig, f"{dataset_name}_sparsity")

    return {"stratified": records, "detection_gamma_r": float(r_det)}


# ============================================================================
# 6. CI COVERAGE BREAKDOWN
# ============================================================================

def analysis_ci_breakdown():
    """Examine where DeepPTR CI coverage fails on synthetic data."""
    print(f"\n{'=' * 60}")
    print("6. CI COVERAGE BREAKDOWN (synthetic)")
    print("=" * 60)

    from scptr.deep.synthetic import generate_kinetic_data

    adata, truth = generate_kinetic_data(n_cells=1500, n_genes=100, seed=0)

    torch.set_num_threads(4)
    model, history = scptr.deep.fit_deepptr(
        adata, d_T=8, d_PT=8, d_hidden=48, n_enc_layers=2,
        batch_size=256, max_epochs=150, kl_warmup_epochs=20,
        patience=15, n_posterior_samples=30,
        device="cpu", seed=0, verbose=False,
    )

    gamma_true = truth["gamma"]
    gamma_mean = adata.layers["gamma"]
    gamma_var = adata.layers["gamma_var"]

    z = 1.96  # 95% CI
    std = np.sqrt(np.clip(gamma_var, 1e-10, None))
    lower = gamma_mean - z * std
    upper = gamma_mean + z * std
    inside = (gamma_true >= lower) & (gamma_true <= upper)

    overall_coverage = float(inside.mean())
    print(f"  Overall 95% CI coverage: {overall_coverage:.4f} (target: 0.95)")

    # Per-gene coverage
    per_gene_cov = inside.mean(axis=0)
    # Per-cell coverage
    per_cell_cov = inside.mean(axis=1)

    # What predicts poor coverage?
    # 1. Genes with high true gamma variance?
    gene_gamma_std = gamma_true.std(axis=0)
    r_cov_std, _ = stats.spearmanr(per_gene_cov, gene_gamma_std)
    print(f"  Per-gene coverage vs true gamma std: r={r_cov_std:.4f}")

    # 2. Coverage by gamma magnitude
    gene_gamma_mean = gamma_true.mean(axis=0)
    r_cov_mean, _ = stats.spearmanr(per_gene_cov, gene_gamma_mean)
    print(f"  Per-gene coverage vs true gamma mean: r={r_cov_mean:.4f}")

    # 3. Is the problem overconfidence (too narrow CI) or bias (wrong mean)?
    error = gamma_mean - gamma_true
    relative_error = np.abs(error) / (gamma_true + 1e-8)
    mean_rel_error = np.median(relative_error)
    mean_ci_width = np.median(2 * z * std)
    mean_true_range = np.median(np.ptp(gamma_true, axis=0))

    print(f"\n  Diagnosis:")
    print(f"    Median relative error: {mean_rel_error:.4f}")
    print(f"    Median 95% CI width:   {mean_ci_width:.4f}")
    print(f"    Median true range:     {mean_true_range:.4f}")
    print(f"    → CI width / true range = {mean_ci_width / max(mean_true_range, 1e-8):.4f}")
    print(f"    → {'Overconfident (CI too narrow)' if overall_coverage < 0.5 else 'Moderate calibration'}")

    result = {
        "overall_coverage": overall_coverage,
        "target_coverage": 0.95,
        "per_gene_cov_vs_std_r": float(r_cov_std),
        "per_gene_cov_vs_mean_r": float(r_cov_mean),
        "median_relative_error": float(mean_rel_error),
        "median_ci_width": float(mean_ci_width),
        "median_true_range": float(mean_true_range),
        "diagnosis": "overconfident" if overall_coverage < 0.5 else "moderate",
    }

    return result


# ============================================================================
# 7. ARE/NMD ENRICHMENT OF PT-SPECIFIC GENES
# ============================================================================

def analysis_pt_gene_enrichment():
    """Are PT-specific genes enriched for ARE or NMD targets?"""
    print(f"\n{'=' * 60}")
    print("7. ARE/NMD ENRICHMENT OF PT-SPECIFIC GENES")
    print("=" * 60)

    are_genes = set()
    with open(DATA_DIR / "are_genes.txt") as f:
        for line in f:
            are_genes.add(line.strip().upper())

    nmd_genes = set()
    with open(DATA_DIR / "nmd_genes.txt") as f:
        for line in f:
            nmd_genes.add(line.strip().upper())

    results = {}
    for dataset_name in ("pancreas", "dentate_gyrus"):
        adv_file = Path(__file__).parent.parent / "output" / "deep_advantages" / "results" / f"{dataset_name}_advantages.json"
        if not adv_file.exists():
            continue

        with open(adv_file) as f:
            adv = json.load(f)

        pt_genes = adv.get("disentanglement", {}).get("pt_specific_genes", [])
        pt_upper = set(g.upper() for g in pt_genes)

        are_overlap = pt_upper & are_genes
        nmd_overlap = pt_upper & nmd_genes

        print(f"\n  {dataset_name}: {len(pt_genes)} PT-specific genes")
        print(f"    ARE overlap: {len(are_overlap)} ({len(are_overlap)/max(len(pt_upper),1)*100:.0f}%)")
        if are_overlap:
            print(f"      {sorted(are_overlap)}")
        print(f"    NMD overlap: {len(nmd_overlap)} ({len(nmd_overlap)/max(len(pt_upper),1)*100:.0f}%)")
        if nmd_overlap:
            print(f"      {sorted(nmd_overlap)}")

        results[dataset_name] = {
            "n_pt_genes": len(pt_genes),
            "are_overlap": sorted(are_overlap),
            "nmd_overlap": sorted(nmd_overlap),
        }

    return results


# ============================================================================
# 8. HONEST LIMITATIONS TABLE
# ============================================================================

def print_limitations():
    print(f"\n{'=' * 60}")
    print("8. HONEST LIMITATIONS")
    print("=" * 60)

    limitations = [
        ("Steady-state assumption", "Violated in actively differentiating cells; dynamic mode requires velocity (circular)"),
        ("Smoothing pre-processing", "Neighbor averaging collapses per-cell variation before gamma estimation"),
        ("Beta estimation", "Upper-quantile regression is crude; beta errors propagate directly into gamma"),
        ("Half-life correlations", "r=-0.35 to -0.40 explains ~15% of variance; modest biological signal"),
        ("sci-fate tautology", "gamma ∝ new/old ≈ ground truth; high correlation is partially structural"),
        ("DeepPTR CI coverage", "27% for 95% CI; posterior is severely overconfident (amortized VI gap)"),
        ("Gene subset", "DeepPTR evaluated on 300 genes for CPU tractability; not full genome"),
        ("No method comparison", "No benchmarking against velVI, DeepVelo, scVI, or other deep methods"),
        ("Single seed", "No error bars; results may vary across random initializations"),
        ("PT-specific genes", "No external perturbation validation; could be technical artifacts"),
        ("Scalability", "Tested on 3K-7K cells; untested on modern 100K+ cell atlases"),
    ]

    for name, desc in limitations:
        print(f"  {name:<25} {desc}")

    return limitations


# ============================================================================
# MAIN
# ============================================================================

def main():
    set_figure_style()
    ensure_dirs()

    all_results = {}

    # Prepare datasets
    datasets = [
        ("pancreas", scptr.datasets.pancreas, "clusters"),
        ("dentate_gyrus", scptr.datasets.dentate_gyrus, "clusters"),
    ]

    for name, loader, cluster_key in datasets:
        print(f"\n{'#' * 60}")
        print(f"# {name.upper()}")
        print(f"{'#' * 60}")

        adata_an = prepare_analytical(loader)
        top_genes = select_top_genes(adata_an, n_top=300)
        ds_results = {}

        # 1. Fair comparison
        ds_results["fair_comparison"] = analysis_fair_comparison(adata_an, name, top_genes)

        # 2. Bootstrap CIs
        ds_results["bootstrap_ci"] = analysis_bootstrap_ci(adata_an, name)

        # 3. eCLIP validation
        ds_results["eclip_validation"] = analysis_eclip_validation(name)

        # 5. Sparsity
        ds_results["sparsity"] = analysis_sparsity(adata_an, name)

        all_results[name] = ds_results

    # 4. sci-fate tautology
    all_results["scifate_tautology"] = analysis_scifate_tautology()

    # 6. CI breakdown (synthetic)
    all_results["ci_breakdown"] = analysis_ci_breakdown()

    # 7. PT gene enrichment
    all_results["pt_enrichment"] = analysis_pt_gene_enrichment()

    # 8. Limitations
    limitations = print_limitations()
    all_results["limitations"] = [{"name": n, "description": d} for n, d in limitations]

    # Save
    with open(OUTPUT_DIR / "results" / "wrapup_results.json", "w") as f:
        json.dump(all_results, f, indent=2, default=str)

    print(f"\n{'=' * 60}")
    print("WRAP-UP COMPLETE")
    print("=" * 60)
    print(f"Results saved to: {OUTPUT_DIR}")


if __name__ == "__main__":
    main()