{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 5. Cell-type Deconvolution\n",
    "\n",
    "**Pipeline phase:** Phase 3.5b\n",
    "\n",
    "**What you will learn:**\n",
    "- Understand the deconvolution module API\n",
    "- Prepare a reference scRNA-seq AnnData\n",
    "- Extract reference profiles with `extract_reference_profiles`\n",
    "- Call `deconvolution` with a supported backend\n",
    "- Inspect `obsm['cell_type_proportions']` and `obs['{method}_argmax']`\n",
    "\n",
    "**Backends supported:** `cell2location`, `tangram`, `destvi`\n",
    "\n",
    "**Install:** `pip install sc-tools[deconvolution]`\n",
    "\n",
    "> **Note:** This tutorial shows the API interface and expected data shapes.\n",
    "> Running a full deconvolution requires a GPU-enabled machine with the\n",
    "> optional `[deconvolution]` extras installed.  The code cells below\n",
    "> are illustrative; they will raise `ImportError` without those extras."
   ]
  },
  {
   "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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare spatial AnnData (target)\n",
    "\n",
    "The spatial AnnData must have raw counts in `adata.X` and spatial coordinates\n",
    "in `adata.obsm['spatial']`.  The `library_id` column in `adata.obs` is used\n",
    "for per-library batching to avoid OOM on large datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(99)\n",
    "n_spots, n_genes = 300, 200\n",
    "gene_names = [f\"GENE_{i}\" for i in range(n_genes)]\n",
    "\n",
    "row, col = np.divmod(np.arange(n_spots), 20)\n",
    "spatial_coords = np.column_stack([col * 100.0, row * 100.0])\n",
    "\n",
    "spatial_adata = ad.AnnData(\n",
    "    X=np.random.poisson(5, size=(n_spots, n_genes)).astype(float),\n",
    "    obs={\"library_id\": [f\"lib_{i % 2}\" for i in range(n_spots)]},\n",
    "    obsm={\"spatial\": spatial_coords},\n",
    ")\n",
    "spatial_adata.var_names = gene_names\n",
    "print(\"Spatial AnnData:\", spatial_adata)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare reference scRNA-seq AnnData\n",
    "\n",
    "The reference must have:\n",
    "- Raw counts in `adata.X`\n",
    "- Cell-type labels in `adata.obs['celltype']`\n",
    "- Overlapping gene names with the spatial AnnData"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_cells = 500\n",
    "cell_types = [\"Epithelial\", \"Macrophage\", \"T_cell\", \"Fibroblast\", \"Endothelial\"]\n",
    "\n",
    "ref_adata = ad.AnnData(\n",
    "    X=np.random.poisson(3, size=(n_cells, n_genes)).astype(float),\n",
    "    obs={\"celltype\": pd.Categorical([cell_types[i % len(cell_types)] for i in range(n_cells)])},\n",
    ")\n",
    "ref_adata.var_names = gene_names\n",
    "print(\"Reference AnnData:\", ref_adata)\n",
    "print(\"Cell types:\", ref_adata.obs[\"celltype\"].value_counts().to_dict())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extract reference profiles\n",
    "\n",
    "`extract_reference_profiles` computes per-cell-type mean expression profiles\n",
    "optimized for Cell2location (negative binomial regression model).\n",
    "It handles memory-efficient batch processing for large reference datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# In a real run (with scvi-tools installed):\n",
    "# inf_aver = tl.extract_reference_profiles(\n",
    "#     ref_adata,\n",
    "#     celltype_col=\"celltype\",\n",
    "#     batch_key=None,\n",
    "#     use_gpu=False,   # set True on GPU machine\n",
    "# )\n",
    "# print(\"Reference profiles shape:\", inf_aver.shape)\n",
    "\n",
    "# Simulated output shape for illustration:\n",
    "print(\"Expected inf_aver shape: (n_genes, n_celltypes) =\", (n_genes, len(cell_types)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run deconvolution\n",
    "\n",
    "The `deconvolution` function accepts `backend=\"cell2location\"`, `\"tangram\"`, or `\"destvi\"`.\n",
    "It processes each `library_id` separately to avoid OOM on large spatial datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# In a real run:\n",
    "# tl.deconvolution(\n",
    "#     spatial_adata,\n",
    "#     reference=ref_adata,\n",
    "#     celltype_col=\"celltype\",\n",
    "#     backend=\"cell2location\",\n",
    "#     library_key=\"library_id\",\n",
    "#     use_gpu=False,\n",
    "#     max_epochs=100,\n",
    "# )\n",
    "\n",
    "# After deconvolution, the following are populated:\n",
    "# - spatial_adata.obsm['cell_type_proportions']  -- shape (n_spots, n_celltypes)\n",
    "# - spatial_adata.obs['cell2location_argmax']     -- dominant cell type per spot\n",
    "\n",
    "print(\"Output: obsm['cell_type_proportions'] shape = (n_spots, n_celltypes)\")\n",
    "print(\"Output: obs['cell2location_argmax'] = dominant cell type per spot\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inspect proportions (synthetic example)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simulate deconvolution output for illustration\n",
    "proportions = np.random.dirichlet(np.ones(len(cell_types)), size=n_spots)\n",
    "spatial_adata.obsm[\"cell_type_proportions\"] = pd.DataFrame(\n",
    "    proportions, index=spatial_adata.obs_names, columns=cell_types\n",
    ")\n",
    "spatial_adata.obs[\"cell2location_argmax\"] = pd.Categorical(\n",
    "    spatial_adata.obsm[\"cell_type_proportions\"].idxmax(axis=1)\n",
    ")\n",
    "\n",
    "print(\"Proportions (first 3 spots):\")\n",
    "print(spatial_adata.obsm[\"cell_type_proportions\"].head(3).round(3).to_string())\n",
    "\n",
    "print(\"\\nDominant cell type distribution:\")\n",
    "print(spatial_adata.obs[\"cell2location_argmax\"].value_counts())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Save and downstream use\n",
    "\n",
    "```python\n",
    "# Save method-specific checkpoint\n",
    "spatial_adata.write_h5ad(\"results/adata.deconvolution.cell2location.h5ad\")\n",
    "\n",
    "# Spatial plots (Phase 5)\n",
    "import sc_tools.pl as pl\n",
    "pl.spatial.plot_spatial_continuous(\n",
    "    spatial_adata,\n",
    "    key=\"cell_type_proportions\",\n",
    "    column=\"Macrophage\",\n",
    "    library_id=\"lib_0\",\n",
    ")\n",
    "```\n",
    "\n",
    "For project-specific deconvolution scripts see:\n",
    "- `projects/visium/ggo_visium/scripts/run_deconvolution.py`\n",
    "- `projects/visium_hd/robin/scripts/run_deconvolution.py`"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "sc_tools",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}