{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Optoinhibition\n",
    "\n",
    "## Purpose\n",
    "Read the optogenetic NWB files and consolidate their trial-level metadata into one `optomat` payload for downstream figures.\n",
    "\n",
    "## What This Notebook Does\n",
    "- Locate all NWB files in `data_optogenetics/`.\n",
    "- Load each trial table using the NWB reader shared by the Python pipeline.\n",
    "- Extract the session identifier, fiber location, and all trial-level columns needed downstream.\n",
    "- Preserve the NWB trial-column order so the exported payload stays close to the source files.\n",
    "- Save the consolidated sessions to `Optoinhibition.npz`.\n",
    "\n",
    "## Inputs\n",
    "- `data_optogenetics/*.nwb`\n",
    "\n",
    "## Outputs\n",
    "- `processed_data/Optoinhibition.npz`\n",
    "\n",
    "## Notebook Flow\n",
    "1. Setup\n",
    "2. Parameters\n",
    "3. Load Data\n",
    "4. Compute\n",
    "5. Preview\n",
    "6. Save\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Setup\n",
    "\n",
    "Resolve the repository root, import the shared NWB loaders, and define the input and output paths used by the optogenetic-session export.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "\n",
    "import gc\n",
    "import sys\n",
    "from pathlib import Path\n",
    "from typing import Any, Dict, List\n",
    "\n",
    "import numpy as np\n",
    "from pynwb import NWBHDF5IO\n",
    "\n",
    "# Resolve the repository root before importing 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.loading_base import to_numpy_compact_series\n",
    "from functions.loading_nwb import extract_fiber_location_from_trials, get_nwb_trials_dataframe\n",
    "from functions.paths import build_default_optoinhibition_paths, resolve_repo_root\n",
    "\n",
    "repo_root = resolve_repo_root(repo_root_hint)\n",
    "paths = build_default_optoinhibition_paths(repo_root)\n",
    "\n",
    "data_dir = paths['data_dir']\n",
    "out_dir = paths['out_dir']\n",
    "out_path = paths['out_path']\n",
    "\n",
    "if not data_dir.exists():\n",
    "    raise FileNotFoundError(\n",
    "        f'Missing directory: {data_dir}. '\n",
    "        'Create/populate data_optogenetics first, then rerun.'\n",
    "    )\n",
    "\n",
    "print(f'Data directory: {data_dir}')\n",
    "print(f'Output file: {out_path.name}')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Parameters\n",
    "\n",
    "Define the minimal set of trial-table columns that must be present in every NWB file before the session can be exported.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "required_trial_fields = {\n",
    "    'start_time',\n",
    "    'stop_time',\n",
    "    'lick_flag',\n",
    "    'lick_time',\n",
    "    'early_lick',\n",
    "    'trial_type',\n",
    "    'opto_window',\n",
    "}\n",
    "\n",
    "np.random.seed(0)\n",
    "print('Required trial fields:', ', '.join(sorted(required_trial_fields)))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Load Data\n",
    "\n",
    "Discover the NWB files that will contribute to the exported `optomat` payload.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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'NWB files found: {len(nwb_files)}')\n",
    "for nwb_path in nwb_files:\n",
    "    print(f'  - {nwb_path.name}')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Compute\n",
    "\n",
    "Read each NWB file, extract the trial table, and build one exported session dictionary per file.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optomat: List[Dict[str, Any]] = []\n",
    "\n",
    "for i_file, nwb_path in enumerate(nwb_files, start=1):\n",
    "    with NWBHDF5IO(str(nwb_path), mode='r', load_namespaces=True) as io:\n",
    "        nwb = io.read()\n",
    "        trials_df = get_nwb_trials_dataframe(nwb)\n",
    "\n",
    "        # Keep the complete trial table columns in their native NWB order so\n",
    "        # the exported payload remains aligned with the source file schema.\n",
    "        columnname_trial = list(trials_df.columns)\n",
    "\n",
    "        missing_required = sorted(required_trial_fields.difference(columnname_trial))\n",
    "        if missing_required:\n",
    "            raise KeyError(\n",
    "                f\"{nwb_path.name}: missing required trial columns: {missing_required}\"\n",
    "            )\n",
    "\n",
    "        sess: Dict[str, Any] = {}\n",
    "        sess['session_id'] = str(getattr(nwb, 'identifier', nwb_path.stem))\n",
    "        sess['fiber_location'] = extract_fiber_location_from_trials(trials_df, opto_area_column='opto_area')\n",
    "\n",
    "        for col in columnname_trial:\n",
    "            sess[col] = to_numpy_compact_series(trials_df[col])\n",
    "\n",
    "        optomat.append(sess)\n",
    "\n",
    "    print(f'[{i_file}/{len(nwb_files)}] {nwb_path.name} -> columns={len(columnname_trial)}')\n",
    "\n",
    "    del trials_df, sess\n",
    "    gc.collect()\n",
    "\n",
    "print(f'Built optomat sessions: {len(optomat)}')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Preview\n",
    "\n",
    "Inspect one exported session dictionary before writing the full payload to disk.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if len(optomat) > 0:\n",
    "    example = optomat[0]\n",
    "    print('Example keys:', list(example.keys())[:20])\n",
    "    print('session_id:', example.get('session_id', ''))\n",
    "    print('fiber_location:', example.get('fiber_location', ''))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Save\n",
    "\n",
    "Write the consolidated optogenetic-session payload to `processed_data/Optoinhibition.npz`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optomat_obj = np.array(optomat, dtype=object)\n",
    "\n",
    "np.savez_compressed(\n",
    "    out_path,\n",
    "    optomat=optomat_obj,\n",
    ")\n",
    "\n",
    "print(f'Saved optoinhibition payload: {out_path.name}')\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
}