scrna_explorer.py

"""
Single-Cell AnnData Explorer — Corrected & Upgraded
=====================================================
Run with:  streamlit run scrna_explorer.py

Fixes applied vs original:
  Bug #1  — @st.cache_data + bytes input (was @st.cache_resource, broken on re-upload)
  Bug #2  — vectorised gene expression lookup (was O(n²) get_loc loop)
  Bug #3  — umap.reindex() instead of .loc[] (was KeyError on filtered indices)
  Bug #4  — safe matrix slice bounds (was crash on small datasets)
  Warn #1 — dynamic categorical column detection for UMAP colors
  Warn #2 — improved matrix type heuristic (max value + sparsity)
  Warn #3 — removed 500-gene cap in Gene Explorer
  Warn #4 — raw/layers fallback for gene expression values
  Add     — multi-QC violin plots, dot plot, export .h5ad, doublet summary
"""

import io
import os
import sys
import tempfile
import warnings

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
try:
    from scrna_report import build_report
    REPORT_AVAILABLE = True
except ImportError:
    REPORT_AVAILABLE = False
import sys
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
try:
    from scrna_report import build_report
    REPORT_AVAILABLE = True
except ImportError:
    REPORT_AVAILABLE = False

import numpy as np
import pandas as pd
import scanpy as sc
import streamlit as st
import plotly.express as px
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

warnings.filterwarnings("ignore")

# ── Raise Streamlit upload limit programmatically ─────────────────────────────
# The default Streamlit limit is 2 GB. For large .h5ad files we override it
# here AND via .streamlit/config.toml (place that file next to this script).
# The environment variable is the only way to set it before the server starts
# when running without a config file.
os.environ.setdefault("STREAMLIT_SERVER_MAX_UPLOAD_SIZE", "10240")  # 10 GB

# ── Page config ───────────────────────────────────────────────────────────────
st.set_page_config(
    page_title="scRNA-seq AnnData Explorer",
    page_icon="🔬",
    layout="wide",
    initial_sidebar_state="expanded",
)

st.markdown("""
<style>
  .metric-card{background:#f8f9fa;border-radius:10px;padding:1rem;text-align:center;border:1px solid #e9ecef}
  .metric-val {font-size:1.9rem;font-weight:600;color:#212529}
  .metric-lbl {font-size:0.76rem;color:#6c757d;margin-top:4px}
  .fix-box   {background:#e8f4fd;border-left:4px solid #1a73e8;padding:.7rem 1rem;
               border-radius:0 8px 8px 0;margin:.75rem 0;font-size:.85rem;color:#1a3a5c}
  h1{font-size:1.55rem!important}
  h2{font-size:1.2rem!important}
  h3{font-size:1rem!important}
</style>
""", unsafe_allow_html=True)


# ── Helpers ───────────────────────────────────────────────────────────────────
def safe_dense(x):
    """Convert sparse or dense matrix slice to a dense numpy array."""
    return x.toarray() if hasattr(x, "toarray") else np.asarray(x)


def get_expression(adata, gene_name: str) -> np.ndarray:
    """
    FIX Bug #2 + Warn #4:
    Vectorised single-gene extraction. Prefers adata.raw when available
    (raw counts are more interpretable for teaching), falls back to adata.X.
    Never uses a per-cell get_loc loop.
    """
    source = adata.raw if adata.raw is not None else adata
    gene_idx = np.where(source.var_names == gene_name)[0]
    if len(gene_idx) == 0:
        return np.zeros(adata.n_obs)
    expr = safe_dense(source.X[:, gene_idx[0]]).flatten()
    return expr


def infer_matrix_status(adata) -> dict:
    """
    FIX Bug #4 + Warn #2:
    Safe slice bounds + improved heuristic using max value and sparsity
    rather than the unreliable all-integers check.
    """
    n_cells = min(1000, adata.n_obs)
    n_genes = min(100, adata.n_vars)
    x = safe_dense(adata.X[:n_cells, :n_genes])

    x_max   = float(np.max(x))
    x_min   = float(np.min(x))
    sparsity = float(np.mean(x == 0))
    # Raw count matrices: large integer-valued max, high sparsity
    likely_raw = (x_max > 20) and (sparsity > 0.5) and np.all(np.mod(x[x > 0], 1) == 0)
    return {
        "min": round(x_min, 4),
        "max": round(x_max, 4),
        "dtype": str(x.dtype),
        "sparsity": f"{sparsity*100:.1f}%",
        "likely_raw": likely_raw,
        "interpretation": (
            "Likely raw / count matrix (high sparsity, integer values, large max)"
            if likely_raw else
            "Likely normalized / log-transformed (non-integer or low max values)"
        ),
    }


def categorical_columns(df: pd.DataFrame) -> list:
    """
    FIX Warn #1:
    Return all columns that are object/category/bool dtype — i.e. suitable
    for coloring a UMAP scatter. No hardcoded column name list.
    """
    return [
        c for c in df.columns
        if df[c].dtype.name in ("object", "category", "bool")
        or (df[c].nunique() <= 50 and df[c].dtype.name != "float64")
    ]


def numeric_columns(df: pd.DataFrame) -> list:
    """Return columns with numeric dtype."""
    return df.select_dtypes(include=[np.number]).columns.tolist()


def export_adata_button(adata, label="Download processed .h5ad", key="dl_h5ad"):
    """
    FIX Add: Export filtered/processed AnnData as a downloadable .h5ad file.
    Writes to a BytesIO buffer via a temp file.
    """
    with tempfile.NamedTemporaryFile(suffix=".h5ad", delete=False) as tmp:
        adata.write_h5ad(tmp.name)
        tmp.seek(0)
        with open(tmp.name, "rb") as f:
            data = f.read()
    st.download_button(label, data=data, file_name="processed_subset.h5ad",
                       mime="application/octet-stream", key=key)


# ── Load data (FIX Bug #1) ────────────────────────────────────────────────────
@st.cache_data(show_spinner="Reading .h5ad file...")
def load_adata(file_bytes: bytes):
    """
    FIX Bug #1: Accept raw bytes, not the UploadedFile object.
    @st.cache_data hashes the bytes correctly, so re-uploading a different
    file correctly invalidates the cache. @st.cache_resource did not.
    """
    return sc.read_h5ad(io.BytesIO(file_bytes))


# ── App header ────────────────────────────────────────────────────────────────
st.title("Single-Cell AnnData Explorer")
st.caption("Explore metadata, QC, UMAP, annotations, and gene expression from your .h5ad dataset.")

uploaded_file = st.file_uploader(
    "Upload your .h5ad file",
    type=["h5ad"],
    help=(
        "Upload limit raised to 10 GB via .streamlit/config.toml. "
        "If the UI still shows '2 GB', ensure config.toml is in a "
        ".streamlit/ folder next to this script and restart Streamlit. "
        "For 10x Genomics data, convert first: sc.read_10x_mtx() → adata.write_h5ad()."
    ),
)

if uploaded_file is None:
    st.info("Upload a .h5ad file to get started.")
    st.stop()

file_mb = uploaded_file.size / 1e6
if file_mb > 2000:
    st.info(
        f"Large file detected ({file_mb:,.0f} MB). "
        "Reading into memory may take 1–3 minutes — please wait and do not refresh."
    )

# FIX Bug #1 — pass bytes, not the file object
adata = load_adata(uploaded_file.read())

obs = adata.obs.copy()
var = adata.var.copy()
var["gene_name"] = var.index.astype(str)
matrix_info = infer_matrix_status(adata)

# ── Sidebar ───────────────────────────────────────────────────────────────────
with st.sidebar:
    st.markdown("## Filters")

    # Sample column selection
    possible_sample_cols = [c for c in obs.columns if "sample" in c.lower()]
    fallback_col = possible_sample_cols[0] if possible_sample_cols else obs.columns[0]
    default_col  = "Sample name" if "Sample name" in obs.columns else fallback_col

    sample_col = st.selectbox(
        "Sample column",
        possible_sample_cols if possible_sample_cols else obs.columns.tolist(),
        index=(possible_sample_cols.index(default_col)
               if default_col in possible_sample_cols else 0),
    )

    all_samples = sorted(obs[sample_col].astype(str).unique().tolist())
    selected_samples = st.multiselect("Select sample(s)", all_samples, default=all_samples)

    filtered_obs = obs[obs[sample_col].astype(str).isin(selected_samples)].copy()

    # Dynamic extra filters — any low-cardinality categorical column
    extra_cat_cols = [
        c for c in categorical_columns(filtered_obs)
        if c != sample_col and filtered_obs[c].nunique() <= 100
    ]
    for col in extra_cat_cols[:6]:   # cap at 6 to avoid sidebar overflow
        vals = sorted(filtered_obs[col].astype(str).dropna().unique().tolist())
        sel  = st.multiselect(f"Filter: {col}", vals, default=vals, key=f"filt_{col}")
        filtered_obs = filtered_obs[filtered_obs[col].astype(str).isin(sel)]

    st.divider()
    st.markdown("### Teaching mode")
    show_tips = st.toggle("Show teaching notes", value=True)


def tip(text: str):
    if show_tips:
        st.markdown(f'<div class="fix-box">💡 {text}</div>', unsafe_allow_html=True)


# ── Tabs ──────────────────────────────────────────────────────────────────────
tab1, tab2, tab3, tab4, tab5, tab6, tab7, tab8 = st.tabs([
    "Overview", "Metadata", "Sample summary",
    "QC dashboard", "Annotation", "UMAP", "Gene explorer", "Report",
])


# ═════════════════════════════════════════════════════════════════════════════
# TAB 1 — OVERVIEW
# ═════════════════════════════════════════════════════════════════════════════
with tab1:
    st.subheader("Dataset overview")

    tip("The AnnData object (adata) is the core data structure in Scanpy. "
        "adata.X is the expression matrix (cells × genes). adata.obs stores "
        "per-cell metadata. adata.var stores per-gene metadata. adata.obsm "
        "stores embeddings like PCA and UMAP.")

    # Summary metrics
    c1, c2, c3, c4, c5 = st.columns(5)
    def mcard(col, val, lbl, color="#212529"):
        col.markdown(
            f'<div class="metric-card"><div class="metric-val" style="color:{color}">'
            f'{val}</div><div class="metric-lbl">{lbl}</div></div>',
            unsafe_allow_html=True,
        )

    mcard(c1, f"{adata.n_obs:,}",       "Total cells")
    mcard(c2, f"{filtered_obs.shape[0]:,}", "Filtered cells", "#1D9E75")
    mcard(c3, f"{adata.n_vars:,}",       "Genes")
    mcard(c4, f"{len(obs.columns)}",     "Metadata columns")
    mcard(c5, "Yes" if "X_umap" in adata.obsm else "No", "UMAP available",
          "#1D9E75" if "X_umap" in adata.obsm else "#D85A30")

    # FIX Add: doublet summary in overview
    if "predicted_doublets" in obs.columns:
        st.markdown("---")
        n_doublets = int(obs["predicted_doublets"].sum()) if obs["predicted_doublets"].dtype == bool \
                     else int((obs["predicted_doublets"].astype(str) == "True").sum())
        pct = n_doublets / adata.n_obs * 100
        dc1, dc2 = st.columns(2)
        mcard(dc1, f"{n_doublets:,}", "Predicted doublets",
              "#D85A30" if pct > 10 else "#212529")
        mcard(dc2, f"{pct:.1f}%", "Doublet rate",
              "#D85A30" if pct > 10 else "#212529")

    st.markdown("---")
    st.markdown("### AnnData structure")
    col_l, col_r = st.columns(2)
    with col_l:
        st.markdown(f"**obs columns ({len(obs.columns)}):**")
        st.code(", ".join(obs.columns.tolist()), language=None)
        st.markdown(f"**var columns ({len(var.columns)}):**")
        st.code(", ".join(var.columns.tolist()), language=None)
    with col_r:
        st.markdown("**obsm keys:**")
        st.code(", ".join(adata.obsm.keys()) or "none", language=None)
        st.markdown("**layers:**")
        st.code(", ".join(adata.layers.keys()) or "none", language=None)
        st.markdown(f"**has raw:** {adata.raw is not None}")

    st.markdown("---")
    st.markdown("### Matrix status")
    mi = matrix_info
    mc1, mc2, mc3, mc4 = st.columns(4)
    mcard(mc1, mi["dtype"],     "dtype")
    mcard(mc2, mi["min"],       "min value")
    mcard(mc3, mi["max"],       "max value")
    mcard(mc4, mi["sparsity"],  "sparsity")
    fn = st.info if not mi["likely_raw"] else st.warning
    fn(f"Interpretation: {mi['interpretation']}")

    # Dual sample column check
    if "Sample name" in obs.columns and "sample name" in obs.columns:
        same = (obs["Sample name"].astype(str) == obs["sample name"].astype(str)).all()
        st.markdown("---")
        st.markdown("### Sample column check")
        dc1, dc2, dc3 = st.columns(3)
        mcard(dc1, obs["Sample name"].nunique(), "'Sample name' unique values")
        mcard(dc2, obs["sample name"].nunique(), "'sample name' unique values")
        mcard(dc3, "Yes" if same else "No",    "Identical columns?",
              "#1D9E75" if same else "#D85A30")
        if not same:
            st.warning("'Sample name' and 'sample name' are not identical — "
                       "they likely represent different levels of sample identity.")
            compare_df = pd.DataFrame({
                "Sample name": obs["Sample name"].astype(str),
                "sample name": obs["sample name"].astype(str),
            }).drop_duplicates().sort_values("Sample name")
            st.dataframe(compare_df, use_container_width=True)

    st.markdown("---")
    st.markdown("### Expression matrix preview (first 5 cells × 5 genes)")
    X_preview = safe_dense(adata.X[:5, :5])
    st.dataframe(
        pd.DataFrame(X_preview, index=adata.obs_names[:5], columns=adata.var_names[:5]),
        use_container_width=True,
    )


# ═════════════════════════════════════════════════════════════════════════════
# TAB 2 — METADATA TABLE
# ═════════════════════════════════════════════════════════════════════════════
with tab2:
    st.subheader("Cell metadata table")

    tip("adata.obs is a pandas DataFrame. Each row is a cell (barcode). "
        "Columns can be QC metrics, experimental conditions, or cluster labels.")

    default_cols = list(filtered_obs.columns[:10])
    selected_cols = st.multiselect(
        "Columns to display", list(filtered_obs.columns), default=default_cols,
    )
    table_df = (filtered_obs[selected_cols] if selected_cols else filtered_obs).copy()
    table_df = table_df.reset_index().rename(columns={"index": "cell_barcode"})
    st.dataframe(table_df, use_container_width=True)

    csv = table_df.to_csv(index=False).encode("utf-8")
    st.download_button(
        "Download filtered metadata (.csv)", data=csv,
        file_name="filtered_metadata.csv", mime="text/csv",
    )


# ═════════════════════════════════════════════════════════════════════════════
# TAB 3 — SAMPLE SUMMARY
# ═════════════════════════════════════════════════════════════════════════════
with tab3:
    st.subheader("Sample summary")

    sample_summary = (
        filtered_obs.groupby(sample_col).size()
        .reset_index(name="n_cells")
        .sort_values("n_cells", ascending=False)
    )
    st.dataframe(sample_summary, use_container_width=True)
    fig = px.bar(sample_summary, x=sample_col, y="n_cells",
                 title=f"Cells per {sample_col}", color_discrete_sequence=["#1D9E75"])
    st.plotly_chart(fig, use_container_width=True)

    compare_options = [c for c in filtered_obs.columns if c != sample_col]
    if compare_options:
        compare_col = st.selectbox("Cross-tabulate against", compare_options)
        ctab = pd.crosstab(
            filtered_obs[sample_col].astype(str),
            filtered_obs[compare_col].astype(str),
        )
        st.dataframe(ctab, use_container_width=True)

    if "Sample name" in filtered_obs.columns and "sample name" in filtered_obs.columns:
        st.markdown("---")
        st.markdown("### 'Sample name' vs 'sample name' mapping")
        mapping = pd.crosstab(
            filtered_obs["Sample name"].astype(str),
            filtered_obs["sample name"].astype(str),
        )
        st.dataframe(mapping, use_container_width=True)


# ═════════════════════════════════════════════════════════════════════════════
# TAB 4 — QC DASHBOARD
# ═════════════════════════════════════════════════════════════════════════════
with tab4:
    st.subheader("QC dashboard")

    tip("Three key QC metrics: (1) n_genes_by_counts — genes detected per cell; "
        "(2) total_counts — total UMIs (sequencing depth); "
        "(3) pct_counts_mt — mitochondrial %. "
        "Dying cells have very high mito % because cytoplasmic mRNA leaks out "
        "but mitochondrial mRNA stays. Doublets show unusually high gene counts.")

    qc_num_cols = [
        c for c in ["n_genes", "n_genes_by_counts", "total_counts",
                    "total_counts_mt", "pct_counts_mt", "doublet_scores"]
        if c in filtered_obs.columns
    ]

    if not qc_num_cols:
        st.warning("No standard QC columns found (n_genes_by_counts, total_counts, pct_counts_mt). "
                   "Run sc.pp.calculate_qc_metrics() first.")
    else:
        # FIX Add: multi-variable violin plots side by side
        st.markdown("### Violin plots")
        n_cols = min(3, len(qc_num_cols))
        violin_cols = qc_num_cols[:n_cols]
        fig_v, axes = plt.subplots(1, n_cols, figsize=(5 * n_cols, 4))
        if n_cols == 1:
            axes = [axes]
        colors = ["#1D9E75", "#7F77DD", "#D85A30", "#378ADD", "#BA7517", "#D4537E"]
        for ax, col, color in zip(axes, violin_cols, colors):
            data_v = filtered_obs[col].dropna()
            parts  = ax.violinplot(data_v, showmedians=True, widths=0.7)
            parts["bodies"][0].set_facecolor(color + "55")
            parts["bodies"][0].set_edgecolor(color)
            parts["cmedians"].set_color(color)
            ax.set_title(col, fontsize=10, fontweight="500")
            ax.set_xticks([])
            ax.tick_params(labelsize=8)
            ax.spines[["top", "right", "bottom"]].set_visible(False)
        plt.tight_layout()
        st.pyplot(fig_v)
        plt.close()

        # Summary stats table
        st.markdown("### Summary statistics")
        stats_df = filtered_obs[qc_num_cols].describe().T.round(2)
        st.dataframe(stats_df, use_container_width=True)

        # Interactive histogram
        st.markdown("### Distribution histogram")
        qc_plot_col = st.selectbox("Variable", qc_num_cols, key="qc_hist_col")
        n_bins = st.slider("Bins", 20, 100, 50)
        fig_h = px.histogram(
            filtered_obs, x=qc_plot_col, nbins=n_bins,
            title=f"Distribution of {qc_plot_col}",
            color_discrete_sequence=["#1D9E75"],
        )
        st.plotly_chart(fig_h, use_container_width=True)

    # Doublet bar chart
    if "predicted_doublets" in filtered_obs.columns:
        st.markdown("---")
        st.markdown("### Doublet predictions")
        tip("Doublets are droplets that captured two cells. They appear as "
            "outliers with unusually high gene counts. Scrublet / DoubletFinder "
            "flags them computationally before clustering.")
        dc = filtered_obs["predicted_doublets"].astype(str).value_counts().reset_index()
        dc.columns = ["predicted_doublets", "count"]
        fig_d = px.bar(dc, x="predicted_doublets", y="count",
                       title="Predicted doublets", color_discrete_sequence=["#D85A30"])
        st.plotly_chart(fig_d, use_container_width=True)


# ═════════════════════════════════════════════════════════════════════════════
# TAB 5 — ANNOTATION & BIOLOGY
# ═════════════════════════════════════════════════════════════════════════════
with tab5:
    st.subheader("Annotation and biological metadata")

    tip("Cell type annotations are usually stored in adata.obs as a categorical "
        "column (e.g. 'annotation', 'cell_type', 'leiden'). Each cluster is "
        "named after its dominant marker gene expression profile.")

    # FIX Warn #1: dynamic categorical columns, not hardcoded list
    bio_options = categorical_columns(filtered_obs)
    if not bio_options:
        st.info("No categorical metadata columns found in the filtered selection.")
    else:
        bio_col = st.selectbox("Metadata column to summarize", bio_options)
        bio_counts = (
            filtered_obs[bio_col].astype(str)
            .value_counts().reset_index()
        )
        bio_counts.columns = [bio_col, "count"]
        bio_counts["pct"] = (bio_counts["count"] / bio_counts["count"].sum() * 100).round(1)

        st.dataframe(bio_counts, use_container_width=True)

        fig_b = px.bar(bio_counts, x=bio_col, y="count",
                       title=f"Cell counts by {bio_col}",
                       color=bio_col, text="pct",
                       color_discrete_sequence=px.colors.qualitative.Safe)
        fig_b.update_traces(texttemplate="%{text}%", textposition="outside")
        fig_b.update_layout(showlegend=False)
        st.plotly_chart(fig_b, use_container_width=True)

        # FIX Add: dot plot of top marker genes grouped by annotation
        if "leiden" in obs.columns or any("annot" in c.lower() for c in obs.columns):
            st.markdown("---")
            st.markdown("### Marker gene dot plot")
            tip("sc.pl.dotplot shows the fraction of cells expressing a gene (dot size) "
                "and the mean expression level (dot color) per group. "
                "This is the standard way to validate cluster annotations.")

            # Discover marker genes present in adata
            DEFAULT_MARKERS = [
                "LGR5","ASCL2","MKI67","TOP2A","VIL1","FABP1",
                "MUC2","TFF3","CHGA","NEUROG3","DCLK1","DEFA5",
            ]
            available_markers = [g for g in DEFAULT_MARKERS if g in adata.var_names]

            custom_genes = st.text_input(
                "Custom gene list (comma-separated, overrides defaults)",
                placeholder="LGR5, MKI67, VIL1, MUC2",
            )
            if custom_genes:
                gene_list = [g.strip() for g in custom_genes.split(",") if g.strip() in adata.var_names]
                missing   = [g.strip() for g in custom_genes.split(",") if g.strip() not in adata.var_names]
                if missing:
                    st.warning(f"Genes not found in dataset: {', '.join(missing)}")
            else:
                gene_list = available_markers

            groupby_col = st.selectbox(
                "Group by", [c for c in bio_options if c in adata.obs.columns],
                key="dotplot_group",
            )

            if gene_list and groupby_col in adata.obs.columns:
                try:
                    fig_dot, ax_dot = plt.subplots(figsize=(max(8, len(gene_list)), 4))
                    sc.pl.dotplot(adata, gene_list, groupby=groupby_col,
                                  ax=ax_dot, show=False)
                    plt.tight_layout()
                    st.pyplot(fig_dot)
                    plt.close()
                except Exception as e:
                    st.warning(f"Dot plot could not be rendered: {e}")
            elif not gene_list:
                st.info("No valid marker genes found in this dataset for the dot plot.")


# ═════════════════════════════════════════════════════════════════════════════
# TAB 6 — UMAP EXPLORER
# ═════════════════════════════════════════════════════════════════════════════
with tab6:
    st.subheader("UMAP explorer")

    tip("UMAP (Uniform Manifold Approximation and Projection) projects the "
        "high-dimensional expression space into 2D. Cells close together share "
        "similar transcriptional profiles. UMAP is for visualization only — "
        "never use UMAP coordinates for quantitative downstream analyses.")

    if "X_umap" not in adata.obsm:
        st.warning("No UMAP found in adata.obsm. Run sc.tl.umap(adata) first.")
    else:
        umap_df = pd.DataFrame(
            adata.obsm["X_umap"], columns=["UMAP1", "UMAP2"], index=adata.obs_names
        ).join(obs)

        # FIX Bug #3: use reindex instead of .loc to avoid KeyError
        umap_df = umap_df.reindex(filtered_obs.index).dropna(subset=["UMAP1", "UMAP2"])

        # FIX Warn #1: dynamic column detection — categorical and numeric
        cat_cols_umap = [c for c in categorical_columns(umap_df)
                         if c not in ("UMAP1", "UMAP2")]
        num_cols_umap = [c for c in numeric_columns(umap_df)
                         if c not in ("UMAP1", "UMAP2")]
        all_color_opts = cat_cols_umap + num_cols_umap

        if not all_color_opts:
            st.warning("No suitable columns found to color the UMAP.")
        else:
            col_a, col_b = st.columns([3, 1])
            with col_a:
                color_col = st.selectbox("Color by", all_color_opts)
            with col_b:
                point_opacity = st.slider("Opacity", 0.1, 1.0, 0.7, 0.05)

            max_pts = len(umap_df)
            n_plot  = st.slider(
                "Cells to plot",
                min_value=min(1000, max_pts),
                max_value=max_pts,
                value=min(20000, max_pts),
                step=max(500, max_pts // 40),
            )
            umap_plot = umap_df.sample(n=n_plot, random_state=42) if n_plot < max_pts else umap_df

            is_cat = color_col in cat_cols_umap
            fig_u  = px.scatter(
                umap_plot, x="UMAP1", y="UMAP2",
                color=color_col,
                opacity=point_opacity,
                title=f"UMAP — colored by {color_col}",
                color_continuous_scale="Viridis" if not is_cat else None,
                color_discrete_sequence=px.colors.qualitative.Safe if is_cat else None,
            )
            fig_u.update_traces(marker_size=3)
            fig_u.update_layout(legend=dict(itemsizing="constant"))
            st.plotly_chart(fig_u, use_container_width=True)

            if n_plot < max_pts:
                st.caption(f"Showing {n_plot:,} of {max_pts:,} cells (random subsample).")


# ═════════════════════════════════════════════════════════════════════════════
# TAB 7 — GENE EXPLORER
# ═════════════════════════════════════════════════════════════════════════════
with tab7:
    st.subheader("Gene explorer")

    tip("Gene expression values in adata.X may be normalized. When adata.raw "
        "is available, expression values shown here come from the raw counts "
        "layer for interpretability. Use the dot plot in the Annotation tab "
        "for multi-gene comparisons.")

    gene_query = st.text_input("Search gene name (partial match ok)")

    filtered_var = var.copy()
    if gene_query:
        filtered_var = filtered_var[
            filtered_var["gene_name"].str.contains(gene_query, case=False, na=False)
        ]

    st.markdown(f"**{len(filtered_var):,} genes matched**")
    # FIX Warn #3: no 500-gene cap — Streamlit selectbox handles large lists natively
    st.dataframe(filtered_var.reset_index(drop=True), use_container_width=True)

    if filtered_var.empty:
        st.info("No genes match the search. Try a shorter query.")
    else:
        selected_gene = st.selectbox(
            "Select gene for expression view",
            options=filtered_var["gene_name"].tolist(),
        )

        if selected_gene:
            # FIX Bug #2 + Warn #4: vectorised + raw fallback
            expr_all = get_expression(adata, selected_gene)
            source_label = "raw counts" if adata.raw is not None else "adata.X"

            expr_df = filtered_obs.copy()
            # Align by position using obs index order
            all_idx  = pd.Index(adata.obs_names)
            filt_pos = all_idx.get_indexer(filtered_obs.index)
            expr_df["expression"] = expr_all[filt_pos]

            st.markdown(f"### {selected_gene} — expression summary ({source_label})")
            st.dataframe(expr_df[["expression"]].describe().T.round(4),
                         use_container_width=True)

            # Group-by box plot
            # FIX Warn #1: dynamic grouping columns
            group_options = [c for c in categorical_columns(expr_df)
                             if expr_df[c].nunique() <= 60]
            if group_options:
                group_col = st.selectbox("Group expression by", group_options, key="gene_group")
                fig_g = px.box(
                    expr_df, x=group_col, y="expression",
                    title=f"{selected_gene} expression by {group_col}",
                    color=group_col,
                    color_discrete_sequence=px.colors.qualitative.Safe,
                    points="outliers",
                )
                fig_g.update_layout(showlegend=False)
                st.plotly_chart(fig_g, use_container_width=True)

                # Violin alternative
                if st.checkbox("Show as violin plot instead"):
                    fig_vg = px.violin(
                        expr_df, x=group_col, y="expression", box=True,
                        title=f"{selected_gene} violin by {group_col}",
                        color=group_col,
                        color_discrete_sequence=px.colors.qualitative.Safe,
                    )
                    fig_vg.update_layout(showlegend=False)
                    st.plotly_chart(fig_vg, use_container_width=True)

            # UMAP overlay for selected gene
            if "X_umap" in adata.obsm:
                st.markdown("---")
                st.markdown(f"### {selected_gene} expression on UMAP")
                umap_gene = pd.DataFrame(
                    adata.obsm["X_umap"], columns=["UMAP1", "UMAP2"], index=adata.obs_names
                ).reindex(filtered_obs.index).dropna()
                umap_gene["expression"] = expr_all[
                    all_idx.get_indexer(umap_gene.index)
                ]
                max_umap_pts = len(umap_gene)
                n_umap = min(20000, max_umap_pts)
                umap_gene_plot = umap_gene.sample(n=n_umap, random_state=0) \
                                 if n_umap < max_umap_pts else umap_gene

                fig_ug = px.scatter(
                    umap_gene_plot, x="UMAP1", y="UMAP2",
                    color="expression",
                    color_continuous_scale="RdYlBu_r",
                    title=f"{selected_gene} expression projected on UMAP",
                    opacity=0.7,
                )
                fig_ug.update_traces(marker_size=3)
                st.plotly_chart(fig_ug, use_container_width=True)

# ── Footer: export processed subset ──────────────────────────────────────────
st.divider()
st.markdown("### Export")
tip("Exporting a filtered AnnData object preserves all metadata, embeddings, "
    "and layers for your collaborators or for reproducible downstream analysis.")

if st.button("Prepare filtered .h5ad export"):
    filt_adata = adata[filtered_obs.index].copy()
    export_adata_button(filt_adata, label=f"Download filtered subset ({filt_adata.n_obs:,} cells)")

# ═════════════════════════════════════════════════════════════════════════════
# TAB 8 — REPORT GENERATOR
# ═════════════════════════════════════════════════════════════════════════════
with tab8:
    st.subheader("Generate analysis report")

    tip("The report is a self-contained HTML file with all charts embedded as "
        "images. Open it in any browser and use File → Print → Save as PDF "
        "to produce a publication-ready PDF. No extra software required.")

    if not REPORT_AVAILABLE:
        st.error(
            "scrna_report.py not found. Place it in the same folder as this "
            "app and restart Streamlit."
        )
    elif not st.session_state.get("processed", False) and "cluster_summary" not in adata.uns:
        st.info(
            "Run the full pipeline first (scrna_pipeline.py) to generate "
            "cluster annotations and developmental stage scores before "
            "building the report."
        )
    else:
        col_r1, col_r2 = st.columns([2, 1])
        with col_r1:
            dataset_name = st.text_input(
                "Dataset name (shown on report cover)",
                value=uploaded_file.name.replace(".h5ad", "") if uploaded_file else "Organoid dataset",
            )
        with col_r2:
            st.markdown("<br>", unsafe_allow_html=True)
            generate_btn = st.button("Build report", type="primary", use_container_width=True)

        if generate_btn:
            with st.spinner("Rendering charts and building report — this takes ~20 seconds..."):
                try:
                    html_report = build_report(adata, dataset_name=dataset_name)
                    st.success("Report built successfully!")
                    st.download_button(
                        label="Download HTML report",
                        data=html_report.encode("utf-8"),
                        file_name=f"{dataset_name.replace(' ', '_')}_scrna_report.html",
                        mime="text/html",
                        use_container_width=True,
                    )
                    st.caption(
                        "Open the downloaded file in Chrome, Firefox, or Safari. "
                        "Use File → Print → Save as PDF to convert to PDF."
                    )

                    # Preview summary of what's in the report
                    st.markdown("---")
                    st.markdown("**Report contents:**")
                    sections = [
                        ("1", "Dataset summary", f"{adata.n_obs:,} cells · {adata.n_vars:,} genes"),
                        ("2", "Quality control", "Violin plots for genes/cell, UMI counts, mito %"),
                        ("3", "UMAP & clustering", "Embedding coloured by Leiden cluster"),
                        ("4", "Cell type annotation", f"{len(adata.uns.get('cluster_summary', []))} clusters annotated"),
                        ("5", "Developmental stage", "Maturity scores + radar charts per cluster"),
                        ("6", "Methods", "Full pipeline description"),
                        ("7", "References", "5 primary citations"),
                    ]
                    for num, title, detail in sections:
                        st.markdown(
                            f"<div style='display:flex;gap:10px;padding:7px 0;"
                            f"border-bottom:1px solid #f2f2f2;font-size:13px'>"
                            f"<span style='width:20px;height:20px;border-radius:50%;"
                            f"background:#1D9E75;color:white;display:flex;align-items:center;"
                            f"justify-content:center;font-size:10px;font-weight:600;flex-shrink:0'>{num}</span>"
                            f"<span style='font-weight:500'>{title}</span>"
                            f"<span style='color:#888;margin-left:auto'>{detail}</span></div>",
                            unsafe_allow_html=True
                        )

                except Exception as e:
                    st.error(f"Report generation failed: {e}")
                    raise e