{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d215ecbe",
   "metadata": {},
   "source": [
    "# Delay Decoding\n",
    "\n",
    "## Purpose\n",
    "Estimate time-resolved decoding accuracy across the task timeline with one binary classifier per time bin.\n",
    "\n",
    "## What This Notebook Does\n",
    "- Load the processed PSTH entries and area definitions.\n",
    "- Configure the delay-decoding analysis across target areas.\n",
    "- Run the time-resolved decoding and the shuffled-label control.\n",
    "- Save the decoding outputs used downstream.\n",
    "\n",
    "## Inputs\n",
    "- `processed_data/psth_10ms.npz`\n",
    "- `data_helpers/Area_list.npz`\n",
    "\n",
    "## Outputs\n",
    "- `processed_data/Decoding_delay.npz`\n",
    "\n",
    "## Notebook Flow\n",
    "1. Environment and paths\n",
    "2. Analysis configuration\n",
    "3. Load inputs\n",
    "4. Prepare analysis inputs\n",
    "5. Run delay decoding\n",
    "6. Save outputs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3a915f8",
   "metadata": {},
   "source": [
    "## 1. Environment and Paths\n",
    "\n",
    "Set imports, resolve the repository root, and define the input and output paths used by this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ee5279a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "\n",
    "import sys\n",
    "from pathlib import Path\n",
    "from typing import Any, Dict, List\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "np.set_printoptions(precision=4, suppress=True)\n",
    "\n",
    "cwd = Path.cwd().resolve()\n",
    "repo_root = next(\n",
    "    (p for p in [cwd, *cwd.parents] if (p / 'functions').is_dir()),\n",
    "    None,\n",
    ")\n",
    "if repo_root is None:\n",
    "    raise FileNotFoundError('Could not locate repo root containing functions')\n",
    "if str(repo_root) not in sys.path:\n",
    "    sys.path.insert(0, str(repo_root))\n",
    "\n",
    "from functions.constants import FUNCTION_API_VERSION\n",
    "from functions.decoding import run_delay_decoding_pipeline\n",
    "from functions.loading import load_canonical_psth_entries_and_area_list\n",
    "from functions.paths import build_default_psth_area_output_paths\n",
    "\n",
    "paths = build_default_psth_area_output_paths(\n",
    "    repo_root,\n",
    "    out_filename='Decoding_delay.npz',\n",
    "    psth_filename='psth_10ms.npz',\n",
    "    area_list_filename='Area_list.npz',\n",
    ")\n",
    "psth_path = paths['psth_path']\n",
    "area_list_path = paths['area_list_path']\n",
    "out_dir = paths['out_dir']\n",
    "out_path = paths['out_path']\n",
    "\n",
    "if not psth_path.exists():\n",
    "    raise FileNotFoundError(f'File not found: {psth_path}')\n",
    "if not area_list_path.exists():\n",
    "    raise FileNotFoundError(f'File not found: {area_list_path}')\n",
    "\n",
    "print('Environment ready.')\n",
    "print(f'API version: {FUNCTION_API_VERSION}')\n",
    "print(f'Inputs: {psth_path.name}, {area_list_path.name}')\n",
    "print(f'Output: {out_path.name}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27fb2258",
   "metadata": {},
   "source": [
    "## 2. Analysis Configuration\n",
    "\n",
    "Define the trial filters, decoding options, and target areas used by the delay-decoding pipeline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f32c247f",
   "metadata": {},
   "outputs": [],
   "source": [
    "params: Dict[str, Any] = {\n",
    "    'quietstate': 'Quiet_(jaw & whisker)',\n",
    "    'BaselineSubtraction': 0,\n",
    "    'completion_state': 'completed_trials',\n",
    "    'TrialType': [1, 2, 3, 4, 5],\n",
    "    'LickState': [1, 0, 0, 0, 0],\n",
    "    'celltype': 'All',\n",
    "    'regionlist': ['A1', 'wM2', 'wS1', 'ALM', 'wS2'],\n",
    "    'mintrial': 5,\n",
    "    'balance_method': 'downsample',\n",
    "    'zscoring': 1,\n",
    "}\n",
    "\n",
    "print('Configuration ready.')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a7bbfa8",
   "metadata": {},
   "source": [
    "## 3. Load Inputs\n",
    "\n",
    "Load the PSTH entries and area definitions into analysis-ready structures for the downstream decoding pipeline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35163dbd",
   "metadata": {},
   "outputs": [],
   "source": [
    "canon = load_canonical_psth_entries_and_area_list(psth_path, area_list_path)\n",
    "\n",
    "entries_raw = canon['entries_raw']\n",
    "entries_clean = canon['entries_clean']\n",
    "entries_by_area = canon['entries_by_area']\n",
    "area_list = canon['area_list']\n",
    "enable_ccf_filter = True\n",
    "windowCenters = np.asarray(canon['window_centers'], dtype=np.float32)\n",
    "\n",
    "print('Regions:', sorted(area_list.keys()))\n",
    "print('Raw entries:', len(entries_raw))\n",
    "print('Prepared entries:', len(entries_clean))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9402fe51",
   "metadata": {},
   "source": [
    "## 4. Prepare Analysis Inputs\n",
    "\n",
    "Attach the shared time axis and final decoding parameters before running the main pipeline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c057449e",
   "metadata": {},
   "outputs": [],
   "source": [
    "params['windowCenters'] = windowCenters\n",
    "params['min_cells_per_session'] = 5\n",
    "\n",
    "print('Areas found:', sorted(entries_by_area.keys()))\n",
    "print('Prepared entries:', len(entries_clean))\n",
    "print('Time bins:', windowCenters.size)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c04c9441",
   "metadata": {},
   "source": [
    "## 5. Run Delay Decoding\n",
    "\n",
    "Run the area-wise delay-decoding pipeline and collect the decoded accuracy matrices together with their shuffled control."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc82745f",
   "metadata": {},
   "outputs": [],
   "source": [
    "rng_seed = 0\n",
    "\n",
    "result = run_delay_decoding_pipeline(\n",
    "    entries_by_area=entries_by_area,\n",
    "    params=params,\n",
    "    area_list=area_list,\n",
    "    enable_ccf_filter=enable_ccf_filter,\n",
    "    rng_seed=rng_seed,\n",
    ")\n",
    "\n",
    "Accuracy = result['Accuracy']\n",
    "Accuracy_shuffeled = result['Accuracy_shuffeled']\n",
    "windowCenters = result['windowCenters']\n",
    "\n",
    "for currentarea in params['regionlist']:\n",
    "    decoded_rows = int(np.sum(Accuracy['sessionaddress'][currentarea] != None))\n",
    "    print(f'{currentarea}: probes={len(entries_by_area.get(currentarea, []))}, decoded={decoded_rows}')\n",
    "\n",
    "print('Delay decoding complete.')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c266432e",
   "metadata": {},
   "source": [
    "## 6. Save Outputs\n",
    "\n",
    "Save the decoding arrays, session-address metadata, time axis, and runtime parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69065c6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.savez_compressed(\n",
    "    out_path,\n",
    "    Accuracy=np.array([Accuracy], dtype=object),\n",
    "    Accuracy_shuffeled=np.array([Accuracy_shuffeled], dtype=object),\n",
    "    windowCenters=np.asarray(windowCenters, dtype=np.float32),\n",
    "    params=np.array([params], dtype=object),\n",
    ")\n",
    "\n",
    "print(f'Wrote: {out_path.name}')"
   ]
  }
 ],
 "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
}