{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d608b99c",
   "metadata": {},
   "source": [
    "# Modulation\n",
    "\n",
    "## Purpose\n",
    "Compute ROC-based modulation metrics for Go-tone whisker Hit trials across the audio, delay, whisker, and lick windows.\n",
    "\n",
    "## What This Notebook Does\n",
    "- Load the processed PSTH bundle used by the population analyses.\n",
    "- Canonicalize the PSTH entries so each probe payload exposes the same trial and unit fields.\n",
    "- Compare each analysis window against its paired baseline window for every cell.\n",
    "- Save discrimination indices, p-values, and firing-rate differences to `Modulation.npz`.\n",
    "\n",
    "## Inputs\n",
    "- `processed_data/psth_10ms.npz`\n",
    "\n",
    "## Outputs\n",
    "- `processed_data/Modulation.npz`\n",
    "\n",
    "## Notebook Flow\n",
    "1. Setup\n",
    "2. Parameters\n",
    "3. Load data\n",
    "4. Canonicalize inputs\n",
    "5. Compute modulation analysis\n",
    "6. Save\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fcc6bbaf",
   "metadata": {},
   "source": [
    "## 1. Setup\n",
    "\n",
    "Resolve the repository root, import the shared utilities used by the notebook, and define the processed-data paths used by the modulation analysis.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2014434",
   "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",
    "# Resolve the repository root before importing shared modules.\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.core import normalize_raw_entries, to_1d, unwrap_scalar_obj\n",
    "from functions.decoding import get_celltype_mask, get_completion_mask, get_quiet_mask\n",
    "from functions.loading import canonicalize_psth_entries_for_decoding\n",
    "from functions.math_utils import selectivity_index_calculation_np\n",
    "\n",
    "# Define the main input and output locations used by this notebook.\n",
    "psth_path = repo_root / 'processed_data' / 'psth_10ms.npz'\n",
    "out_dir = repo_root / 'processed_data'\n",
    "out_dir.mkdir(parents=True, exist_ok=True)\n",
    "out_path = out_dir / 'Modulation.npz'\n",
    "\n",
    "if not psth_path.exists():\n",
    "    raise FileNotFoundError(f'File not found: {psth_path}')\n",
    "\n",
    "print(f'Input file: {psth_path.name}')\n",
    "print(f'Output file: {out_path.name}')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "335faa46",
   "metadata": {},
   "source": [
    "## 2. Parameters\n",
    "\n",
    "Set the trial-selection filters and the four analysis windows described in `Modulation.m`, together with their paired baseline windows.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a17e3ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "rng_seed = 0\n",
    "\n",
    "params: Dict[str, Any] = {\n",
    "    'QuietState': 'Quiet_(jaw & whisker)',\n",
    "    'BaselineSubtraction': 0,\n",
    "    'completion_state': 'completed_trials',\n",
    "    'TrialType': [1],\n",
    "    'LickState': [1],\n",
    "    'CellType': 'All',\n",
    "    'regionlist': ['A1', 'wS1', 'ALM', 'wM2', 'wS2'],\n",
    "}\n",
    "\n",
    "window_name = ['audio', 'delay', 'whisker', 'lick']\n",
    "\n",
    "# legacy bins converted to 0-based Python indexing\n",
    "windows_list1 = [\n",
    "    np.arange(100, 103, dtype=np.int32),  # 101:103\n",
    "    np.arange(180, 200, dtype=np.int32),  # 181:200\n",
    "    np.arange(200, 203, dtype=np.int32),  # 201:203\n",
    "    np.arange(230, 250, dtype=np.int32),  # 231:250\n",
    "]\n",
    "windows_baseline = [\n",
    "    np.arange(97, 100, dtype=np.int32),   # 98:100\n",
    "    np.arange(80, 100, dtype=np.int32),   # 81:100\n",
    "    np.arange(197, 200, dtype=np.int32),  # 198:200\n",
    "    np.arange(80, 100, dtype=np.int32),   # 81:100\n",
    "]\n",
    "\n",
    "bin_width = np.float32(0.01)\n",
    "\n",
    "print(\n",
    "    f\"Configured modulation analysis | QuietState={params['QuietState']} | \"\n",
    "    f\"CellType={params['CellType']} | windows={', '.join(window_name)}\"\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd2e8a2c",
   "metadata": {},
   "source": [
    "## 3. Load Data\n",
    "\n",
    "Load the processed PSTH bundle that contains the trial-level spike-count tensors for each probe entry.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25fb6547",
   "metadata": {},
   "outputs": [],
   "source": [
    "with np.load(psth_path, allow_pickle=True) as data:\n",
    "    if 'psth_mat' not in data:\n",
    "        raise KeyError(\"NPZ missing key 'psth_mat'\")\n",
    "    raw_arr = np.asarray(data['psth_mat']).reshape(-1)\n",
    "\n",
    "entries_raw = normalize_raw_entries(raw_arr)\n",
    "print(f'Loaded raw PSTH entries: {len(entries_raw)}')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c395e8e",
   "metadata": {},
   "source": [
    "## 4. Canonicalize Inputs\n",
    "\n",
    "Normalize the raw PSTH entries into a consistent Python structure before running the window-by-window ROC analysis.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aebcaafb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build safe placeholders for canonicalizer consistency check.\n",
    "gid_max = 0\n",
    "for e in entries_raw:\n",
    "    raw_gid = e.get('GlobalclusterID', None)\n",
    "    if raw_gid is None:\n",
    "        continue\n",
    "    gid = to_1d(unwrap_scalar_obj(raw_gid), np.int64)\n",
    "    if gid.size:\n",
    "        gid_max = max(gid_max, int(np.nanmax(gid)))\n",
    "\n",
    "if gid_max <= 0:\n",
    "    total_units = 0\n",
    "    for e in entries_raw:\n",
    "        sp = np.asarray(unwrap_scalar_obj(e['spike_counts']))\n",
    "        if sp.ndim == 3:\n",
    "            total_units += int(sp.shape[2])\n",
    "    gid_max = max(total_units, 1)\n",
    "\n",
    "id_ordered = np.arange(1, gid_max + 1, dtype=np.int64)\n",
    "neuron_counter = id_ordered.copy()\n",
    "\n",
    "canon = canonicalize_psth_entries_for_decoding(\n",
    "    entries_raw=entries_raw,\n",
    "    id_ordered=id_ordered,\n",
    "    neuron_counter=neuron_counter,\n",
    ")\n",
    "\n",
    "entries_clean = canon['entries_clean']\n",
    "entries_by_area = canon['entries_by_area']\n",
    "\n",
    "if not entries_clean:\n",
    "    raise RuntimeError('No canonical PSTH entries available.')\n",
    "\n",
    "n_bins_ref = entries_clean[0]['spike_counts'].shape[0]\n",
    "for w1, wb in zip(windows_list1, windows_baseline):\n",
    "    if int(np.max(w1)) >= n_bins_ref or int(np.max(wb)) >= n_bins_ref:\n",
    "        raise ValueError('Window bins exceed available PSTH bins')\n",
    "\n",
    "print(f'Canonical PSTH entries: {len(entries_clean)}')\n",
    "print(f\"Areas with entries: {', '.join(sorted(entries_by_area.keys()))}\")\n",
    "print(f'Reference PSTH bins: {n_bins_ref}')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea839129",
   "metadata": {},
   "source": [
    "## 5. Compute\n",
    "\n",
    "For each brain area and behavioral window, compare activity in the target window against its paired baseline window and store the resulting modulation metrics.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55fc3100",
   "metadata": {},
   "outputs": [],
   "source": [
    "rng = np.random.default_rng(rng_seed)\n",
    "\n",
    "modulation: Dict[str, Any] = {}\n",
    "\n",
    "total_iterations = len(window_name) * len(params['regionlist'])\n",
    "iteration_counter = 0\n",
    "\n",
    "for ind_window, curr_win_name in enumerate(window_name):\n",
    "    roc_mat: List[Any] = [None] * len(entries_clean)\n",
    "\n",
    "    for current_area in params['regionlist']:\n",
    "        probe_entries = list(entries_by_area.get(current_area, []))\n",
    "\n",
    "        for e in probe_entries:\n",
    "            ind_probe = int(e['entry_idx'])\n",
    "            sp_cnt_condition: List[np.ndarray] = []\n",
    "\n",
    "            for ind_cond in range(len(params['TrialType'])):\n",
    "                trial = e['trial']\n",
    "                lick = e['lick']\n",
    "\n",
    "                ind_trial_type = trial == params['TrialType'][ind_cond]\n",
    "                ind_lick_state = lick == params['LickState'][ind_cond]\n",
    "\n",
    "                completion_state = get_completion_mask(e, params['completion_state'])\n",
    "                qind = get_quiet_mask(e, params['QuietState'])\n",
    "                # Apply the behavioral filters before extracting the condition-specific spike tensor.\n",
    "                current_trial_ind = qind & completion_state & ind_lick_state & ind_trial_type\n",
    "\n",
    "                celltype_ind = get_celltype_mask(e, params['CellType'])\n",
    "                # Modulation.m does not apply an additional CCF filter at this stage.\n",
    "                current_cell_ind = celltype_ind\n",
    "\n",
    "                curr_sp = e['spike_counts']\n",
    "                curr_sp_trials = curr_sp[:, current_trial_ind, :]\n",
    "                curr_sp_trials_cells = curr_sp_trials[:, :, current_cell_ind]\n",
    "                sp_cnt_condition.append(curr_sp_trials_cells)\n",
    "\n",
    "            if not sp_cnt_condition or sp_cnt_condition[0].ndim != 3:\n",
    "                continue\n",
    "\n",
    "            cond_data = sp_cnt_condition[0]\n",
    "            n_trials = cond_data.shape[1]\n",
    "            n_cells = cond_data.shape[2]\n",
    "\n",
    "            if n_trials == 0 or n_cells == 0:\n",
    "                roc_mat[ind_probe] = {\n",
    "                    'diff_fr': np.zeros((0, 1), dtype=np.float32),\n",
    "                    'discrimination_index': np.zeros((0, 1), dtype=np.float32),\n",
    "                    'pvalue': np.zeros((0, 1), dtype=np.float32),\n",
    "                }\n",
    "                continue\n",
    "\n",
    "            # Use the MATLAB-defined response and baseline windows for the current modulation period.\n",
    "            bins_range1 = windows_list1[ind_window]\n",
    "            bins_baseline = windows_baseline[ind_window]\n",
    "\n",
    "            p_value = np.full((n_cells, 1), np.nan, dtype=np.float32)\n",
    "            discrimination_index = np.full((n_cells, 1), np.nan, dtype=np.float32)\n",
    "            diff_fr = np.full((n_cells, 1), np.nan, dtype=np.float32)\n",
    "\n",
    "            # Compare each cell's activity in the response window against its paired baseline window.\n",
    "            for ind_cells in range(n_cells):\n",
    "                s_analysis = np.sum(cond_data[bins_range1, :, ind_cells], axis=0)\n",
    "                s_base = np.sum(cond_data[bins_baseline, :, ind_cells], axis=0)\n",
    "                sp_cnt_bin_cell = np.concatenate([s_analysis, s_base]).astype(np.float32, copy=False)\n",
    "\n",
    "                label = np.concatenate([\n",
    "                    np.ones((n_trials,), dtype=np.int16),\n",
    "                    2 * np.ones((n_trials,), dtype=np.int16),\n",
    "                ])\n",
    "\n",
    "                di, p, _auc = selectivity_index_calculation_np(\n",
    "                    sp_cnt_bin_cell,\n",
    "                    label,\n",
    "                    method='permut',\n",
    "                    permutations=200,\n",
    "                    nboot=200,\n",
    "                    rng_obj=rng,\n",
    "                    pos_label=1,\n",
    "                )\n",
    "\n",
    "                discrimination_index[ind_cells, 0] = np.float32(di)\n",
    "                p_value[ind_cells, 0] = np.float32(p)\n",
    "\n",
    "                fr_analysis = np.mean(cond_data[bins_range1, :, ind_cells] / bin_width)\n",
    "                fr_base = np.mean(cond_data[bins_baseline, :, ind_cells] / bin_width)\n",
    "                diff_fr[ind_cells, 0] = np.float32(fr_analysis - fr_base)\n",
    "\n",
    "            roc_mat[ind_probe] = {\n",
    "                'diff_fr': diff_fr,\n",
    "                'discrimination_index': discrimination_index,\n",
    "                'pvalue': p_value,\n",
    "            }\n",
    "\n",
    "        iteration_counter += 1\n",
    "        print(f'Progress {iteration_counter}/{total_iterations} | window={curr_win_name} | area={current_area}')\n",
    "\n",
    "    modulation[curr_win_name] = {'roc_mat': np.array(roc_mat, dtype=object)}\n",
    "\n",
    "print(f'Finished modulation analysis across {len(window_name)} windows.')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b63f0225",
   "metadata": {},
   "source": [
    "## 6. Save\n",
    "\n",
    "Write the modulation payload to `processed_data/Modulation.npz` for reuse in downstream notebooks.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "220ebfb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.savez_compressed(\n",
    "    out_path,\n",
    "    modulation=np.array([modulation], dtype=object),\n",
    "    params=np.array([params], dtype=object),\n",
    "    window_name=np.array(window_name, dtype=object),\n",
    ")\n",
    "\n",
    "print(f'Saved modulation 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
}