{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 6. Spatial Analysis\n",
    "\n",
    "**Pipeline phase:** Phase 5 — Downstream Biology\n",
    "\n",
    "**What you will learn:**\n",
    "- Compute co-localization scores (`truncated_similarity`, `pearson_correlation`)\n",
    "- Run Moran's I spatial autocorrelation per library\n",
    "- Run neighborhood enrichment analysis\n",
    "- Generate a multipage spatial PDF with signature overlays\n",
    "\n",
    "**Prerequisites:** `adata.normalized.scored.p35.h5ad` (Phase 3.5b checkpoint)"
   ]
  },
  {
   "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.tl.colocalization import (\n",
    "    truncated_similarity,\n",
    "    pearson_correlation,\n",
    "    morans_i_batch,\n",
    "    neighborhood_enrichment_batch,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Synthetic spatial AnnData with signature scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(42)\n",
    "n_spots, n_genes = 400, 100\n",
    "\n",
    "# Spatial grid\n",
    "row, col = np.divmod(np.arange(n_spots), 20)\n",
    "spatial = np.column_stack([col * 100.0, row * 100.0])\n",
    "\n",
    "gene_names = [f\"GENE_{i}\" for i in range(n_genes)]\n",
    "X = np.random.rand(n_spots, n_genes).astype(np.float32)\n",
    "\n",
    "adata = ad.AnnData(X=X, obsm={\"spatial\": spatial})\n",
    "adata.var_names = gene_names\n",
    "adata.obs[\"library_id\"] = pd.Categorical([f\"lib_{i % 2}\" for i in range(n_spots)])\n",
    "adata.obs[\"leiden\"] = pd.Categorical([str(i % 5) for i in range(n_spots)])\n",
    "\n",
    "# Simulate signature scores in obsm (as produced by score_signature)\n",
    "sig_cols = [\"Hallmark/HYPOXIA\", \"Hallmark/MYC_TARGETS_V1\", \"Myeloid/Macrophage_Core\"]\n",
    "score_data = np.random.randn(n_spots, len(sig_cols)).astype(np.float32)\n",
    "# Make HYPOXIA and Macrophage positively correlated in one region\n",
    "score_data[:100, 0] += 1.5\n",
    "score_data[:100, 2] += 1.2\n",
    "\n",
    "adata.obsm[\"signature_score\"] = pd.DataFrame(score_data, index=adata.obs_names, columns=sig_cols)\n",
    "adata.obsm[\"signature_score_z\"] = adata.obsm[\"signature_score\"].apply(lambda c: (c - c.mean()) / c.std())\n",
    "\n",
    "print(adata)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Co-localization: truncated similarity\n",
    "\n",
    "`truncated_similarity` computes `score_a * score_b` where both are positive,\n",
    "else 0.  This highlights spots where two programs are jointly active."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hypoxia_scores = adata.obsm[\"signature_score\"][\"Hallmark/HYPOXIA\"].values\n",
    "macro_scores   = adata.obsm[\"signature_score\"][\"Myeloid/Macrophage_Core\"].values\n",
    "\n",
    "coloc = truncated_similarity(hypoxia_scores, macro_scores)\n",
    "print(f\"Co-localization score: mean={coloc.mean():.3f}, max={coloc.max():.3f}\")\n",
    "print(f\"Fraction of spots with joint signal: {(coloc > 0).mean():.2%}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Co-localization: Pearson correlation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sc_tools.tl.colocalization import pearson_correlation\n",
    "\n",
    "corr_df = pearson_correlation(\n",
    "    adata,\n",
    "    sig_a=\"Hallmark/HYPOXIA\",\n",
    "    sig_b=\"Myeloid/Macrophage_Core\",\n",
    "    groupby=\"library_id\",\n",
    ")\n",
    "print(corr_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Moran's I spatial autocorrelation (per library)\n",
    "\n",
    "`morans_i_batch` runs Moran's I for each signature column separately per library.\n",
    "Requires squidpy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# In a real run (squidpy must be installed):\n",
    "# morans_results = morans_i_batch(\n",
    "#     adata,\n",
    "#     keys=sig_cols,\n",
    "#     library_key=\"library_id\",\n",
    "#     n_perms=100,\n",
    "# )\n",
    "# print(morans_results)\n",
    "print(\"morans_i_batch returns a DataFrame with columns: key, library_id, I, pval, fdr\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Neighborhood enrichment\n",
    "\n",
    "`neighborhood_enrichment_batch` tests whether cell types co-occur in spatial\n",
    "neighborhoods more than expected by chance (squidpy `gr.nhood_enrichment` wrapper)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# In a real run:\n",
    "# nhood_df = neighborhood_enrichment_batch(\n",
    "#     adata,\n",
    "#     cluster_key=\"leiden\",\n",
    "#     library_key=\"library_id\",\n",
    "# )\n",
    "# print(nhood_df)\n",
    "print(\"neighborhood_enrichment_batch returns a long-format DataFrame:\")\n",
    "print(\"  columns: library_id, celltype_A, celltype_B, zscore, pval, fdr\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Multipage spatial PDF\n",
    "\n",
    "`multipage_spatial_pdf` generates a PDF with one page per signature per library,\n",
    "overlaid on the tissue slide."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tempfile, os\n",
    "\n",
    "with tempfile.TemporaryDirectory() as tmpdir:\n",
    "    pdf_path = os.path.join(tmpdir, \"spatial_signatures.pdf\")\n",
    "    pl.spatial.multipage_spatial_pdf(\n",
    "        adata,\n",
    "        keys=sig_cols,\n",
    "        library_key=\"library_id\",\n",
    "        output_path=pdf_path,\n",
    "        spot_size=50,\n",
    "    )\n",
    "    print(f\"PDF written: {os.path.getsize(pdf_path)} bytes\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Next steps\n",
    "\n",
    "For a real project, Phase 5 scripts live under `projects/<platform>/<project>/scripts/`.\n",
    "Example:\n",
    "\n",
    "```bash\n",
    "# Run spatial multipage colocalization plot for ggo_visium\n",
    "make -C projects/visium/ggo_visium spatial-multipage-colocalization\n",
    "```\n",
    "\n",
    "Publication figures are saved to `projects/<platform>/<project>/figures/manuscript/`."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "sc_tools",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}