{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualization Tutorial\n",
    "\n",
    "This tutorial covers NeuRodent's comprehensive visualization capabilities for EEG analysis results.\n",
    "\n",
    "## Overview\n",
    "\n",
    "NeuRodent provides two main plotting classes:\n",
    "\n",
    "1. **AnimalPlotter**: Visualize data from a single animal\n",
    "2. **ExperimentPlotter**: Compare data across multiple animals with grouping\n",
    "\n",
    "Both support various plot types: time series, categorical plots, heatmaps, and more."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "from pathlib import Path\n",
    "import logging\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "\n",
    "from neurodent import core, visualization, constants\n",
    "\n",
    "logging.basicConfig(level=logging.INFO)\n",
    "logger = logging.getLogger()\n",
    "\n",
    "# Set plotting style\n",
    "sns.set_style(\"whitegrid\")\n",
    "plt.rcParams['figure.figsize'] = (12, 6)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Single Animal Visualization (AnimalPlotter)\n",
    "\n",
    "### Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load WAR for a single animal\n",
    "war_path = Path(\"/path/to/war/animal_001\")\n",
    "war = visualization.WindowAnalysisResult.load_pickle_and_json(war_path)\n",
    "\n",
    "# Create plotter\n",
    "ap = visualization.AnimalPlotter(war)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Time Series Plots\n",
    "\n",
    "Visualize features over time:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot RMS amplitude over time\n",
    "fig = ap.plot_feature_over_time('rms')\n",
    "plt.title(\"RMS Amplitude Over Time\")\n",
    "plt.show()\n",
    "\n",
    "# Plot log RMS\n",
    "fig = ap.plot_feature_over_time('logrms')\n",
    "plt.title(\"Log RMS Amplitude Over Time\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Band Power Visualization\n",
    "\n",
    "Display power across frequency bands:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot band powers\n",
    "fig = ap.plot_psdband()\n",
    "plt.title(\"Power Spectral Density by Band\")\n",
    "plt.show()\n",
    "\n",
    "# Plot band powers over time\n",
    "fig = ap.plot_psdband_over_time()\n",
    "plt.title(\"Band Power Over Time\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Multi-Animal Comparison (ExperimentPlotter)\n",
    "\n",
    "### Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load multiple WARs\n",
    "war_paths = [\n",
    "    Path(\"/path/to/war/animal_001\"),\n",
    "    Path(\"/path/to/war/animal_002\"),\n",
    "    Path(\"/path/to/war/animal_003\"),\n",
    "]\n",
    "\n",
    "wars = [\n",
    "    visualization.WindowAnalysisResult.load_pickle_and_json(path)\n",
    "    for path in war_paths\n",
    "]\n",
    "\n",
    "# Create experiment plotter\n",
    "ep = visualization.ExperimentPlotter(\n",
    "    wars,\n",
    "    exclude=['nspike', 'lognspike']  # Exclude features\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Categorical Plots\n",
    "\n",
    "Compare features across groups using box plots, violin plots, swarm plots, etc."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Box plot grouped by genotype\n",
    "g = ep.plot_catplot(\n",
    "    'rms',\n",
    "    groupby='genotype',\n",
    "    kind='box',\n",
    "    catplot_params={'showfliers': False}\n",
    ")\n",
    "plt.suptitle(\"RMS by Genotype\")\n",
    "plt.show()\n",
    "\n",
    "# Violin plot with day/night comparison\n",
    "g = ep.plot_catplot(\n",
    "    'rms',\n",
    "    groupby=['genotype', 'isday'],\n",
    "    x='genotype',\n",
    "    col='isday',\n",
    "    kind='violin'\n",
    ")\n",
    "plt.suptitle(\"RMS by Genotype and Time of Day\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Band Power Comparisons"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Box plot of band powers by genotype\n",
    "g = ep.plot_catplot(\n",
    "    'psdband',\n",
    "    groupby='genotype',\n",
    "    x='genotype',\n",
    "    hue='band',\n",
    "    kind='box',\n",
    "    collapse_channels=True,\n",
    "    catplot_params={'showfliers': False}\n",
    ")\n",
    "plt.suptitle(\"Band Power by Genotype\")\n",
    "plt.show()\n",
    "\n",
    "# With day/night split\n",
    "g = ep.plot_catplot(\n",
    "    'psdband',\n",
    "    groupby=['genotype', 'isday'],\n",
    "    x='genotype',\n",
    "    col='isday',\n",
    "    hue='band',\n",
    "    kind='box',\n",
    "    collapse_channels=True,\n",
    "    catplot_params={'showfliers': False}\n",
    ")\n",
    "plt.suptitle(\"Band Power by Genotype and Time\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Swarm and Point Plots\n",
    "\n",
    "Show individual data points with averages:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Swarm plot with animal-level averaging\n",
    "g = ep.plot_catplot(\n",
    "    'rms',\n",
    "    groupby=['animal', 'genotype'],\n",
    "    x='genotype',\n",
    "    hue='channel',\n",
    "    kind='swarm',\n",
    "    average_groupby=True,\n",
    "    collapse_channels=False,\n",
    "    catplot_params={'dodge': True, 'errorbar': 'ci'}\n",
    ")\n",
    "plt.suptitle(\"RMS by Genotype (Animal Averages)\")\n",
    "plt.show()\n",
    "\n",
    "# Point plot with confidence intervals\n",
    "g = ep.plot_catplot(\n",
    "    'rms',\n",
    "    groupby=['animal', 'genotype'],\n",
    "    x='genotype',\n",
    "    kind='point',\n",
    "    average_groupby=True,\n",
    "    collapse_channels=True,\n",
    "    catplot_params={'errorbar': 'ci'}\n",
    ")\n",
    "plt.suptitle(\"RMS by Genotype (Mean ± CI)\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Connectivity Visualization\n",
    "\n",
    "### Heatmaps for Coherence and Correlation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Coherence heatmap by genotype\n",
    "g = ep.plot_heatmap(\n",
    "    'cohere',\n",
    "    groupby='genotype'\n",
    ")\n",
    "plt.suptitle(\"Coherence by Genotype\")\n",
    "plt.show()\n",
    "\n",
    "# Pearson correlation heatmap\n",
    "g = ep.plot_heatmap(\n",
    "    'pcorr',\n",
    "    groupby='genotype'\n",
    ")\n",
    "plt.suptitle(\"Correlation by Genotype\")\n",
    "plt.show()\n",
    "\n",
    "# With day/night comparison\n",
    "g = ep.plot_heatmap(\n",
    "    'cohere',\n",
    "    groupby=['genotype', 'isday']\n",
    ")\n",
    "plt.suptitle(\"Coherence by Genotype and Time\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Difference Heatmaps\n",
    "\n",
    "Show differences relative to a baseline condition:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Difference from wildtype\n",
    "g = ep.plot_diffheatmap(\n",
    "    'cohere',\n",
    "    groupby=['genotype', 'isday'],\n",
    "    baseline_key='WT',\n",
    "    baseline_groupby='genotype',\n",
    "    remove_baseline=True\n",
    ")\n",
    "plt.suptitle(\"Coherence Difference from WT\")\n",
    "plt.show()\n",
    "\n",
    "# By frequency band\n",
    "g = ep.plot_diffheatmap(\n",
    "    'cohere',\n",
    "    groupby='genotype',\n",
    "    baseline_key='WT',\n",
    "    baseline_groupby='genotype',\n",
    "    col='band',\n",
    "    row='genotype',\n",
    "    remove_baseline=True\n",
    ")\n",
    "plt.suptitle(\"Band-Specific Coherence Differences\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. QQ Plots for Distribution Analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# QQ plot by genotype and channel\n",
    "g = ep.plot_qqplot(\n",
    "    'rms',\n",
    "    groupby=['genotype'],\n",
    "    row='genotype',\n",
    "    col='channel',\n",
    "    height=3\n",
    ")\n",
    "plt.suptitle(\"RMS Distribution QQ Plot\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Customizing Plots\n",
    "\n",
    "### Plot Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Custom catplot parameters\n",
    "custom_params = {\n",
    "    'showfliers': False,\n",
    "    'aspect': 2,\n",
    "    'height': 5,\n",
    "    'palette': 'Set2'\n",
    "}\n",
    "\n",
    "g = ep.plot_catplot(\n",
    "    'rms',\n",
    "    groupby='genotype',\n",
    "    kind='box',\n",
    "    catplot_params=custom_params\n",
    ")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Channel Collapsing\n",
    "\n",
    "Average across channels:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Without channel collapsing (show all channels)\n",
    "g = ep.plot_catplot(\n",
    "    'rms',\n",
    "    groupby='genotype',\n",
    "    kind='box',\n",
    "    collapse_channels=False\n",
    ")\n",
    "plt.suptitle(\"RMS by Channel\")\n",
    "plt.show()\n",
    "\n",
    "# With channel collapsing (average across channels)\n",
    "g = ep.plot_catplot(\n",
    "    'rms',\n",
    "    groupby='genotype',\n",
    "    kind='box',\n",
    "    collapse_channels=True\n",
    ")\n",
    "plt.suptitle(\"RMS (Averaged Across Channels)\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Saving Figures\n",
    "\n",
    "Save high-quality figures for publications:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create output directory\n",
    "output_folder = Path(\"./figures\")\n",
    "output_folder.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "# Generate and save plot\n",
    "g = ep.plot_catplot(\n",
    "    'rms',\n",
    "    groupby='genotype',\n",
    "    kind='box',\n",
    "    catplot_params={'showfliers': False}\n",
    ")\n",
    "\n",
    "# Save as PNG (high DPI for publications)\n",
    "g.savefig(output_folder / 'rms_by_genotype.png', dpi=300, bbox_inches='tight')\n",
    "\n",
    "# Save as PDF (vector format)\n",
    "g.savefig(output_folder / 'rms_by_genotype.pdf', bbox_inches='tight')\n",
    "\n",
    "print(f\"Figures saved to {output_folder}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Batch Figure Generation\n",
    "\n",
    "Generate multiple figures systematically:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate plots for all linear features\n",
    "for feature in constants.LINEAR_FEATURES:\n",
    "    if feature not in ep.exclude:\n",
    "        logger.info(f\"Generating plot for {feature}\")\n",
    "        \n",
    "        # Box plot\n",
    "        g = ep.plot_catplot(\n",
    "            feature,\n",
    "            groupby='genotype',\n",
    "            kind='box',\n",
    "            collapse_channels=True,\n",
    "            catplot_params={'showfliers': False}\n",
    "        )\n",
    "        g.savefig(\n",
    "            output_folder / f'{feature}_genotype_box.png',\n",
    "            dpi=300,\n",
    "            bbox_inches='tight'\n",
    "        )\n",
    "        plt.close()\n",
    "\n",
    "print(\"Batch generation complete!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Advanced: Custom Plot Types\n",
    "\n",
    "Access underlying data for custom visualizations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get aggregated data from ExperimentPlotter\n",
    "df = ep.get_dataframe('rms', groupby='genotype')\n",
    "\n",
    "# Create custom plot with matplotlib\n",
    "fig, ax = plt.subplots(figsize=(10, 6))\n",
    "\n",
    "# Plot using pandas/seaborn directly\n",
    "sns.boxplot(data=df, x='genotype', y='rms', ax=ax)\n",
    "ax.set_title(\"Custom RMS Plot\")\n",
    "ax.set_ylabel(\"RMS Amplitude\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## Summary\n\nThis tutorial covered:\n\n1. Single animal visualization with AnimalPlotter\n2. Multi-animal comparisons with ExperimentPlotter\n3. Categorical plots (box, violin, swarm, point)\n4. Connectivity visualizations (heatmaps)\n5. Difference heatmaps\n6. Distribution analysis (QQ plots)\n7. Plot customization\n8. Saving figures\n9. Batch figure generation\n10. Custom plot types\n\n## Plot Type Reference\n\n### AnimalPlotter Methods\n- `plot_feature_over_time(feature)`: Time series plot\n- `plot_psdband()`: Bar plot of band powers\n- `plot_psdband_over_time()`: Band powers over time\n\n### ExperimentPlotter Methods\n- `plot_catplot(feature, ...)`: Categorical plots (box, violin, swarm, point, strip)\n- `plot_heatmap(feature, ...)`: Connectivity heatmaps\n- `plot_diffheatmap(feature, ...)`: Difference heatmaps\n- `plot_qqplot(feature, ...)`: Distribution QQ plots\n\n## Next Steps\n\n- **[Windowed Analysis Tutorial](windowed_analysis.ipynb)**: Generate data for visualization\n- **[Basic Usage Tutorial](../quickstart/basic_usage.ipynb)**: Complete workflow"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}