{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial 3 — Visualisation, GSEA, and export\n",
    "\n",
    "This tutorial covers:\n",
    "\n",
    "1. Similarity heatmap\n",
    "2. MP score distributions across experimental groups (violin plots)\n",
    "3. MP scores on UMAP\n",
    "4. Gene Set Enrichment Analysis (GSEA) using bundled MSigDB gene sets\n",
    "5. **Exporting the gene signatures to CSV** — with and without loading scores\n",
    "\n",
    "**Dataset** — Peter Sims lab DMSO vs. panobinostat scRNA-seq cohort."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0 · Imports and data preparation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import kaleidocell\n",
    "import scanpy as sc\n",
    "import pandas as pd\n",
    "\n",
    "adata = sc.read_h5ad('../data/petersims_example.h5ad')\n",
    "\n",
    "results_nmf, nmf_convergence = kaleidocell.multi_sample_nmf(\n",
    "    adata,\n",
    "    batch_key='Patients',\n",
    "    test_ranks=[4, 5, 6, 7, 8, 9],\n",
    "    n_initializations=10,\n",
    "    seed=42,\n",
    ")\n",
    "\n",
    "results_mp = kaleidocell.derive_nmf_metaprograms(results_nmf, top_n_genes=50)\n",
    "mp_scores = kaleidocell.compute_mp_scores(results_mp, adata)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1 · Similarity heatmap\n",
    "\n",
    "Visualise pairwise cosine similarities between all NMF programs.  Diagonal\n",
    "blocks correspond to meta-programs; tight blocks indicate well-defined\n",
    "signatures."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kaleidocell.plot_heatmap(results_mp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2 · Quality metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results_mp['metrics']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3 · MP score distributions across conditions\n",
    "\n",
    "`show_distribution_over_obs` draws one violin plot per MP, grouped by the\n",
    "specified `obs` column.  This is useful for quickly identifying MPs that\n",
    "are differentially active between conditions or donors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kaleidocell.show_distribution_over_obs(\n",
    "    mp_scores,\n",
    "    adata,\n",
    "    batch_key='Treatment',\n",
    "    save=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4 · MP scores on UMAP\n",
    "\n",
    "Overlay per-cell MP scores on the UMAP embedding.  If no UMAP is present in\n",
    "`adata.obsm`, it is computed automatically.\n",
    "\n",
    "Setting `weighted=True` scales each gene's expression by its normalised\n",
    "loading weight before summing, giving higher-specificity genes more\n",
    "influence than low-weight genes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Standard equal-weight scores\n",
    "kaleidocell.plot_mp_scores_on_umap(mp_scores, adata, ncols=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Weight-scaled scores — requires results_mp\n",
    "kaleidocell.plot_mp_scores_on_umap(\n",
    "    mp_scores, adata,\n",
    "    ncols=4,\n",
    "    weighted=True,\n",
    "    results_mp=results_mp,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5 · Gene Set Enrichment Analysis\n",
    "\n",
    "kaleidocell ships with several MSigDB gene-set files.  Use `kaleidocell.files` to\n",
    "discover what is available."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# List all bundled gene-set files with descriptions\n",
    "print(kaleidocell.files)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run Enrichr-based GSEA against Hallmarks and GO Biological Process\n",
    "# Pass the short filename — kaleidocell resolves it to the bundled path automatically\n",
    "gsea_results, gsea_plots = kaleidocell.run_gsea_pipeline(\n",
    "    results_mp,\n",
    "    from_file=[\n",
    "        'h.all.v2026.1.Hs.symbols.gmt',\n",
    "        'c5.go.bp.v2026.1.Hs.symbols.gmt',\n",
    "    ],\n",
    "    top_n_plot=6,\n",
    "    plot=False,\n",
    "    save_csv=False,\n",
    ")\n",
    "\n",
    "# Significant terms table\n",
    "gsea_results.head(20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Bar plots of top terms per MP\n",
    "kaleidocell.plot_gsea_results(gsea_plots, ncols=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6 · Export gene signatures to CSV\n",
    "\n",
    "The gene signatures live in `results_mp['mp_dict']` — a dictionary mapping\n",
    "each MP name to a `pd.Series` of gene names (index) and loading scores\n",
    "(values), sorted descending by score.\n",
    "\n",
    "Two export formats are shown below."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6a · Gene names only (no scores)\n",
    "\n",
    "Long-format table: one row per gene per MP.  Useful for downstream tools\n",
    "that expect a plain gene list per program."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "genes_only = pd.DataFrame(\n",
    "    [\n",
    "        {\"MP\": mp_name, \"gene\": gene}\n",
    "        for mp_name, gene_series in results_mp[\"mp_dict\"].items()\n",
    "        for gene in gene_series.index\n",
    "    ]\n",
    ")\n",
    "\n",
    "genes_only.to_csv('../results/03/metaprograms_genes.csv', index=False)\n",
    "print(f\"Saved {len(genes_only)} rows. Preview:\")\n",
    "genes_only.head(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6b · Gene names with loading scores\n",
    "\n",
    "Same long-format table but with an additional `score` column containing the\n",
    "consensus NMF loading weight.  Higher scores indicate greater specificity\n",
    "to this meta-program."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "genes_with_scores = pd.DataFrame(\n",
    "    [\n",
    "        {\"MP\": mp_name, \"gene\": gene, \"score\": float(score)}\n",
    "        for mp_name, gene_series in results_mp[\"mp_dict\"].items()\n",
    "        for gene, score in gene_series.items()\n",
    "    ]\n",
    ")\n",
    "\n",
    "genes_with_scores.to_csv('../results/03/metaprograms_genes_scores.csv', index=False)\n",
    "print(f\"Saved {len(genes_with_scores)} rows. Preview:\")\n",
    "genes_with_scores.head(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6c · Wide format — one column per MP\n",
    "\n",
    "Useful when you want to open the file in Excel and compare signatures\n",
    "side-by-side."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wide = pd.DataFrame(\n",
    "    {\n",
    "        mp_name: gene_series.index.tolist()\n",
    "        for mp_name, gene_series in results_mp[\"mp_dict\"].items()\n",
    "    }\n",
    ")\n",
    "\n",
    "wide.to_csv('../results/03/metaprograms_wide.csv', index=False)\n",
    "wide.head()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "kaleidocell_07",
   "language": "python",
   "name": "kaleidocell_07"
  },
  "language_info": {
   "name": "python",
   "version": "3.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}