{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4. Gene Signature Scoring\n",
    "\n",
    "**Pipeline phase:** Phase 3.5b\n",
    "\n",
    "**What you will learn:**\n",
    "- Load the bundled MSigDB Hallmark gene sets (offline)\n",
    "- Merge project-specific signatures with Hallmark\n",
    "- Score signatures using `score_signature` (scanpy, UCell, ssGSEA methods)\n",
    "- Read back scores from `adata.obsm`\n",
    "- Run over-representation analysis (ORA)\n",
    "- Plot enrichment results with `plot_gsea_dotplot`\n",
    "\n",
    "**Storage convention:**\n",
    "Scores are written to `adata.obsm['signature_score']` (raw) and\n",
    "`adata.obsm['signature_score_z']` (z-scored).  Column names are full paths,\n",
    "e.g. `Hallmark/HYPOXIA`, `Myeloid/Macrophage_Core`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import anndata as ad\n",
    "import sc_tools.tl as tl\n",
    "import sc_tools.pl as pl\n",
    "from sc_tools.utils.signatures import get_signature_columns, get_signature_df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Synthetic data with known gene sets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(1)\n",
    "n_obs, n_genes = 200, 500\n",
    "\n",
    "# Include some real gene symbols so Hallmark can match them\n",
    "hallmark_sample_genes = [\n",
    "    \"ALDOA\", \"CDKN3\", \"ENO1\", \"LDHA\", \"PGK1\",  # HYPOXIA\n",
    "    \"VEGFA\", \"HK1\", \"HK2\", \"PFKL\", \"TPI1\",\n",
    "    \"MYC\", \"CCND1\", \"CDK4\", \"E2F1\", \"PCNA\",    # MYC_TARGETS_V1\n",
    "    \"TP53\", \"CDKN1A\", \"BBC3\", \"PUMA\", \"BAX\",   # P53_PATHWAY\n",
    "]\n",
    "filler_genes = [f\"GENE_{i}\" for i in range(n_genes - len(hallmark_sample_genes))]\n",
    "gene_names = hallmark_sample_genes + filler_genes\n",
    "\n",
    "X = np.random.rand(n_obs, n_genes).astype(np.float32)\n",
    "\n",
    "adata = ad.AnnData(X=X)\n",
    "adata.var_names = gene_names\n",
    "adata.obs[\"leiden\"] = pd.Categorical([str(i % 4) for i in range(n_obs)])\n",
    "adata.obs[\"condition\"] = pd.Categorical([\"tumor\" if i % 3 == 0 else \"normal\" for i in range(n_obs)])\n",
    "\n",
    "print(adata)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load gene sets\n",
    "\n",
    "### Bundled MSigDB Hallmark (offline)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hallmark = tl.load_hallmark()\n",
    "print(f\"Loaded {len(hallmark)} Hallmark gene sets\")\n",
    "print(\"Example (HYPOXIA):\", list(hallmark[\"Hallmark/HYPOXIA\"])[:8])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Project-specific signatures\n",
    "\n",
    "Project signatures are typically stored in `metadata/gene_signatures.json`.\n",
    "Here we define a minimal example inline:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "project_sigs = {\n",
    "    \"Myeloid/Macrophage_Core\": [\"MRC1\", \"CD68\", \"CSF1R\", \"ITGAM\", \"CD14\"],\n",
    "    \"Myeloid/M2_Polarization\": [\"MRC1\", \"CD163\", \"ARG1\", \"IL10\", \"TGFB1\"],\n",
    "}\n",
    "\n",
    "# Validate and merge with Hallmark\n",
    "report = tl.validate_gene_signatures(project_sigs)\n",
    "print(\"Validation report:\", report)\n",
    "\n",
    "combined = tl.merge_gene_signatures(project_sigs, hallmark)\n",
    "print(f\"Combined: {len(combined)} gene sets\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Score signatures\n",
    "\n",
    "The default method is `\"scanpy\"` (scanpy `tl.score_genes`).\n",
    "Set `method=\"ucell\"` or `method=\"ssgsea\"` for rank-based alternatives\n",
    "(requires `pip install sc-tools[geneset]`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tl.score_signature(adata, combined, method=\"scanpy\")\n",
    "\n",
    "print(\"Scoring method recorded:\", adata.uns.get(\"scoring_method\"))\n",
    "print(\"obsm keys:\", [k for k in adata.obsm if \"signature\" in k])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Retrieve scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# List all scored signatures\n",
    "cols = get_signature_columns(adata)\n",
    "print(f\"Scored {len(cols)} signatures\")\n",
    "print(\"First 5:\", cols[:5])\n",
    "\n",
    "# Get as DataFrame (spots x signatures)\n",
    "score_df = get_signature_df(adata)\n",
    "print(\"\\nScore DataFrame shape:\", score_df.shape)\n",
    "print(score_df.iloc[:3, :3])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Over-representation analysis (ORA)\n",
    "\n",
    "`run_ora` performs Fisher exact test with Benjamini-Hochberg FDR correction\n",
    "for a set of query genes against the loaded gene sets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Query: genes upregulated in cluster 0 (here, synthetic random subset)\n",
    "query_genes = [\"ALDOA\", \"LDHA\", \"PGK1\", \"HK1\", \"VEGFA\", \"ENO1\", \"TPI1\"]\n",
    "\n",
    "ora_results = tl.run_ora(\n",
    "    query_genes=query_genes,\n",
    "    gene_sets=hallmark,\n",
    "    background=adata.var_names.tolist(),\n",
    ")\n",
    "\n",
    "print(ora_results[[\"Term\", \"pval\", \"fdr\", \"overlap\"]].head(5).to_string())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot ORA results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "fig = pl.plot_gsea_dotplot(\n",
    "    ora_results,\n",
    "    fdr_col=\"fdr\",\n",
    "    score_col=\"overlap\",\n",
    "    top_n=10,\n",
    "    title=\"ORA: cluster 0 upregulated genes\",\n",
    ")\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Save Phase 3.5b checkpoint\n",
    "\n",
    "```python\n",
    "adata.write_h5ad(\"results/adata.normalized.scored.p35.h5ad\")\n",
    "```\n",
    "\n",
    "The checkpoint must have:\n",
    "- `obsm['signature_score']` — raw scores (spots x signatures)\n",
    "- `obsm['signature_score_z']` — z-scored scores\n",
    "- `uns['signature_score_report']` — per-signature n_present / n_missing\n",
    "- `uns['scoring_method']` — the method used\n",
    "\n",
    "Downstream scripts read scores via `get_signature_df(adata)` / `get_signature_columns(adata)`."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "sc_tools",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}