{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PSTH 10 ms\n",
    "\n",
    "## Purpose\n",
    "Build the standard-resolution PSTH bundle used by downstream analyses from the raw NWB electrophysiology sessions.\n",
    "\n",
    "## What This Notebook Does\n",
    "- Discover the NWB sessions in `data_electrophysiology/`.\n",
    "- Configure the 10 ms spike and piezo binning used for the export.\n",
    "- Build session-level behavioral and neural PSTH payloads with quiet-trial annotations.\n",
    "- Flatten the session payloads into a probe-level table and assign global unit identifiers.\n",
    "- Save the canonical `psth_10ms.npz` bundle used by the rest of the Python pipeline.\n",
    "\n",
    "## Inputs\n",
    "- `data_electrophysiology/*.nwb`\n",
    "\n",
    "## Outputs\n",
    "- `processed_data/psth_10ms.npz`\n",
    "\n",
    "## Notebook Flow\n",
    "1. Setup\n",
    "2. Parameters\n",
    "3. Load Data\n",
    "4. Compute\n",
    "5. Finalize Outputs\n",
    "6. Save\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Setup\n",
    "\n",
    "Resolve the repository root, import the shared PSTH-building modules, and prepare the processed-data paths used for the 10 ms export.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "from pathlib import Path\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "# Ensure repository root is importable before loading shared modules.\n",
    "cwd = Path.cwd().resolve()\n",
    "repo_root_hint = next(\n",
    "    (p for p in [cwd, *cwd.parents] if (p / 'functions').is_dir()),\n",
    "    None,\n",
    ")\n",
    "if repo_root_hint is None:\n",
    "    raise FileNotFoundError('Could not locate repo root containing functions')\n",
    "if str(repo_root_hint) not in sys.path:\n",
    "    sys.path.insert(0, str(repo_root_hint))\n",
    "\n",
    "from functions.behavior import QuietTrialParams\n",
    "from functions.paths import build_default_ephys_psth_paths, resolve_repo_root\n",
    "from functions.psth_build import (\n",
    "    PsthBuildConfig,\n",
    "    assign_global_cluster_ids,\n",
    "    build_ephys_psth_sessions,\n",
    "    flatten_psth_sessions,\n",
    "    save_psth_entries_npz,\n",
    ")\n",
    "\n",
    "np.set_printoptions(precision=4, suppress=True)\n",
    "repo_root = resolve_repo_root(repo_root_hint)\n",
    "print(f'Repository root: {repo_root}')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Parameters\n",
    "\n",
    "Define the 10 ms PSTH binning and the quiet-trial settings used to build the standard-resolution ephys bundle.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "paths = build_default_ephys_psth_paths(repo_root, out_filename='PSTH_10ms.npz')\n",
    "data_dir = paths['data_dir']\n",
    "out_dir = paths['out_dir']\n",
    "out_path = paths['out_path']\n",
    "legacy_out_path = out_dir / 'PSTH_10ms_structured_full.npz'\n",
    "\n",
    "pre_time = -1.0\n",
    "post_time = 2.0\n",
    "bin_width = 0.01\n",
    "bin_step = 0.01\n",
    "show_progress = True\n",
    "save_legacy_copy = False\n",
    "\n",
    "quiet_params = QuietTrialParams(\n",
    "    prewhisk_window=(0.8, 1.0),\n",
    "    baseline_window=(-1.0, 0.0),\n",
    "    movement_signals=('whisker_speed', 'jaw_movement'),\n",
    "    selection_method='mad_all',\n",
    ")\n",
    "\n",
    "print(\n",
    "    f'Configured PSTH build | pre={pre_time}s | post={post_time}s | '\n",
    "    f'bin_width={bin_width}s | bin_step={bin_step}s'\n",
    ")\n",
    "print(f'Canonical output path: {out_path}')\n",
    "print(f'Legacy-copy enabled: {save_legacy_copy}')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Load Data\n",
    "\n",
    "Inspect the NWB directory and collect the session files that will be converted into the 10 ms PSTH bundle.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not data_dir.exists():\n",
    "    raise FileNotFoundError(f'Missing directory: {data_dir}')\n",
    "\n",
    "nwb_files = sorted(data_dir.glob('*.nwb'))\n",
    "if not nwb_files:\n",
    "    raise FileNotFoundError(f'No .nwb files found in: {data_dir}')\n",
    "\n",
    "print(f'Input NWB directory: {data_dir}')\n",
    "print(f'NWB files discovered: {len(nwb_files)}')\n",
    "print('NWB sessions to process:')\n",
    "for nwb_path in nwb_files:\n",
    "    print(f'  - {nwb_path.name}')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Compute\n",
    "\n",
    "Build one PSTH payload per NWB session using the shared ephys PSTH pipeline.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build one session payload per NWB file before flattening to probe entries.\n",
    "build_cfg = PsthBuildConfig(\n",
    "    pre_time=pre_time,\n",
    "    post_time=post_time,\n",
    "    spike_bin_width=bin_width,\n",
    "    spike_bin_step=bin_step,\n",
    "    piezo_bin_width=bin_width,\n",
    "    piezo_bin_step=bin_step,\n",
    "    quiet_params=quiet_params,\n",
    ")\n",
    "\n",
    "sessions = build_ephys_psth_sessions(nwb_files, build_cfg, show_progress=show_progress)\n",
    "print(f'Session payloads built: {len(sessions)}')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Finalize Outputs\n",
    "\n",
    "Flatten the session payloads into the probe-level export structure and assign global unit identifiers across the full dataset.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psth_mat = flatten_psth_sessions(sessions)\n",
    "offset = assign_global_cluster_ids(psth_mat)\n",
    "\n",
    "print(f'Total probe entries: {len(psth_mat)}')\n",
    "print(f'Total global units: {offset}')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Save\n",
    "\n",
    "Write the canonical 10 ms PSTH bundle to the processed-data directory and emit the optional compatibility copy when requested.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save the canonical NPZ first, then optionally emit the compatibility copy.\n",
    "save_psth_entries_npz(out_path, psth_mat, key='psth_mat')\n",
    "print(f'Saved canonical NPZ: {out_path}')\n",
    "\n",
    "if save_legacy_copy:\n",
    "    save_psth_entries_npz(legacy_out_path, psth_mat, key='psth_mat')\n",
    "    print(f'Saved legacy NPZ: {legacy_out_path}')\n",
    "\n",
    "print('PSTH 10 ms export completed successfully.')\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv (3.9.6)",
   "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.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}