{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PC Projections\n",
    "\n",
    "## Purpose\n",
    "Project single-trial neural trajectories onto the first two principal components to quantify session-wise attractor dynamics.\n",
    "\n",
    "## What This Notebook Does\n",
    "- Load the processed PSTH bundle together with the anatomical area list.\n",
    "- Canonicalize the probe entries so each session exposes the same trial and unit fields.\n",
    "- Select Go-tone whisker lick trials and No-go-tone whisker no-lick trials within each area.\n",
    "- Baseline-subtract each trial, fit a neuron-space PCA from the condition-averaged trajectories, and project all trials onto PC1 and PC2.\n",
    "- Save the session-wise PC trajectories and explained-variance summaries to `PC_projections.npz`.\n",
    "\n",
    "## Inputs\n",
    "- `processed_data/PSTH_10ms.npz`\n",
    "- `data_helpers/Area_list.npz`\n",
    "\n",
    "## Outputs\n",
    "- `processed_data/PC_projections.npz`\n",
    "\n",
    "## Notebook Flow\n",
    "1. Setup\n",
    "2. Parameters\n",
    "3. Load Data\n",
    "4. Canonicalize Inputs\n",
    "5. Compute\n",
    "6. Save\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Setup\n",
    "\n",
    "Resolve the repository root, import the shared loading and selection helpers, and define the processed-data paths used by the PCA projection export.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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",
    "from sklearn.decomposition import PCA\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_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_processed import (\n",
    "    canonicalize_psth_entries_for_decoding,\n",
    "    infer_placeholder_id_order,\n",
    "    load_psth_entries_and_area_list,\n",
    ")\n",
    "from functions.paths import build_default_psth_area_output_paths, resolve_repo_root\n",
    "from functions.selection import get_ccf_mask, get_celltype_mask, get_completion_mask, get_quiet_mask\n",
    "\n",
    "repo_root = resolve_repo_root(repo_root_hint)\n",
    "paths = build_default_psth_area_output_paths(repo_root, out_filename='PC_projections.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",
    "np.random.seed(0)\n",
    "\n",
    "print(f'Input PSTH bundle: {psth_path.name}')\n",
    "print(f'Input area-list bundle: {area_list_path.name}')\n",
    "print(f'Output file: {out_path.name}')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Parameters\n",
    "\n",
    "Set the behavioral filters and PCA settings described in `pc_projections.m`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "params: Dict[str, Any] = {\n",
    "    'QuietState': 'Quiet_(jaw & whisker)',\n",
    "    'completion_state': 'completed_trials',\n",
    "    'TrialType': [1, 3],\n",
    "    'LickState': [1, 0],\n",
    "    'CellType': 'All',\n",
    "    'regionlist': ['A1', 'wS1', 'wS2', 'wM2', 'ALM'],\n",
    "}\n",
    "\n",
    "# Preserve the source-script constants for readability and parity.\n",
    "min_trials_per_session = 5\n",
    "trial_averagedpca = 1\n",
    "\n",
    "print(\n",
    "    f\"Configured PC projections | quiet={params['QuietState']} | \"\n",
    "    f\"celltype={params['CellType']} | areas={', '.join(params['regionlist'])}\"\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Load Data\n",
    "\n",
    "Load the processed PSTH entries and the anatomical area list used for the within-area CCF filter.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "entries_raw, area_list = load_psth_entries_and_area_list(\n",
    "    psth_path=psth_path,\n",
    "    area_list_path=area_list_path,\n",
    ")\n",
    "\n",
    "print(f'Loaded raw PSTH entries: {len(entries_raw)}')\n",
    "print(f\"Area-list regions: {', '.join(sorted(area_list.keys()))}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Canonicalize Inputs\n",
    "\n",
    "Canonicalize the PSTH entries and infer the time-axis convention used by the PCA projections.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "id_ordered, neuron_counter = infer_placeholder_id_order(entries_raw)\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 PSTH entries found after canonicalization.')\n",
    "\n",
    "# Infer the temporal grid from the spike-count tensor length, matching the\n",
    "# source script's 5 ms versus 10 ms branch.\n",
    "sz = int(entries_clean[0]['spike_counts'].shape[0])\n",
    "if sz == 600:\n",
    "    time_axis = np.arange(-1.0 + 0.005, 2.0 + 1e-12, 0.005, dtype=np.float32)\n",
    "    pca_endbin = 600\n",
    "else:\n",
    "    time_axis = np.arange(-1.0 + 0.01, 2.0 + 1e-12, 0.01, dtype=np.float32)\n",
    "    pca_endbin = 300\n",
    "\n",
    "time_axis_full = time_axis.copy()\n",
    "\n",
    "print(f'Canonical PSTH entries: {len(entries_clean)}')\n",
    "print(f\"Areas in PSTH bundle: {', '.join(sorted(entries_by_area.keys()))}\")\n",
    "print(f'Inferred spike-count bins: {sz}')\n",
    "print(f'pca_endbin: {pca_endbin}')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Compute\n",
    "\n",
    "Fit a neuron-space PCA for each valid session, then project every surviving trial onto PC1 and PC2 and store the condition-wise trajectories.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "attractor_results: Dict[str, Any] = {}\n",
    "\n",
    "for current_area in params['regionlist']:\n",
    "    probe_entries = list(entries_by_area.get(current_area, []))\n",
    "    sessions: List[Dict[str, Any]] = []\n",
    "\n",
    "    print(f'Area {current_area}: probes={len(probe_entries)}')\n",
    "\n",
    "    for e in probe_entries:\n",
    "        session_id = str(e.get('session_id', ''))\n",
    "        probe_id = int(e.get('entry_idx', -1))\n",
    "\n",
    "        trial = np.asarray(e['trial']).reshape(-1)\n",
    "        lick = np.asarray(e['lick']).reshape(-1).astype(np.int8)\n",
    "        curr_sc = np.asarray(e['spike_counts'], dtype=np.float32)\n",
    "        window_centers = np.asarray(e['trial_timestamps'], dtype=np.float32)\n",
    "\n",
    "        if curr_sc.ndim != 3 or curr_sc.size == 0:\n",
    "            continue\n",
    "\n",
    "        n_timebins = int(curr_sc.shape[0])\n",
    "        n_trials_total = int(curr_sc.shape[1])\n",
    "\n",
    "        # Trim all trial-level arrays to the common trial count before building masks.\n",
    "        n_trials_common = min(trial.size, lick.size, n_trials_total)\n",
    "        if n_trials_common <= 0:\n",
    "            continue\n",
    "\n",
    "        trial = trial[:n_trials_common]\n",
    "        lick = lick[:n_trials_common]\n",
    "        curr_sc = curr_sc[:, :n_trials_common, :]\n",
    "\n",
    "        # Build a trimmed entry dict so the shared masks operate on the same trial span.\n",
    "        entry = dict(e)\n",
    "        for key in ['trial', 'lick', 'early_lick', 'quiet_whisker', 'quiet_jaw', 'lick_time', 'start_time']:\n",
    "            if key in entry:\n",
    "                entry[key] = np.asarray(entry[key]).reshape(-1)[:n_trials_common]\n",
    "        entry['n_trials'] = n_trials_common\n",
    "\n",
    "        concat_all_trials_data: List[np.ndarray] = []\n",
    "        cond_id: List[np.ndarray] = []\n",
    "        data_for_neuron_pca: List[np.ndarray] = []\n",
    "\n",
    "        baseline_first_bin = int(np.argmin(np.abs(window_centers - (-1.0))))\n",
    "        baseline_last_bin = int(np.argmin(np.abs(window_centers - 0.0)))\n",
    "        if baseline_last_bin < baseline_first_bin:\n",
    "            baseline_first_bin, baseline_last_bin = baseline_last_bin, baseline_first_bin\n",
    "\n",
    "        for icond, trial_type in enumerate(params['TrialType']):\n",
    "            ind_trial_type = trial == int(trial_type)\n",
    "            ind_lick_state = lick == int(params['LickState'][icond])\n",
    "\n",
    "            completion_state = get_completion_mask(entry, params['completion_state'])\n",
    "            qind = get_quiet_mask(entry, params['QuietState'])\n",
    "            curr_trial_ind = qind & completion_state & ind_lick_state & ind_trial_type\n",
    "\n",
    "            celltype_ind = get_celltype_mask(entry, params['CellType'])\n",
    "            ccf_ind = get_ccf_mask(entry, current_area, area_list=area_list, enable_ccf_filter=True)\n",
    "            curr_cell_ind = celltype_ind & ccf_ind\n",
    "\n",
    "            if not np.any(curr_cell_ind):\n",
    "                continue\n",
    "\n",
    "            trial_indices = np.flatnonzero(curr_trial_ind)\n",
    "            n_trials_cond = int(trial_indices.size)\n",
    "            # Code_M pc_projections.m defines min_trials_per_session but never uses it:\n",
    "            # every condition with >=1 trial contributes to the PCA input; empty\n",
    "            # conditions add nothing (matching MATLAB's mean over an empty 3rd dim).\n",
    "            if n_trials_cond < 1:\n",
    "                continue\n",
    "\n",
    "            n_cells_sel = int(np.sum(curr_cell_ind))\n",
    "            all_trials_data = np.zeros((n_cells_sel, n_timebins, n_trials_cond), dtype=np.float32)\n",
    "\n",
    "            # Baseline-subtract each trial before computing the condition mean used by PCA.\n",
    "            for itr, t_idx in enumerate(trial_indices):\n",
    "                curr_sp_trial = curr_sc[:, t_idx, :][:, curr_cell_ind]\n",
    "                baseline_vec = np.nanmean(\n",
    "                    curr_sp_trial[baseline_first_bin:baseline_last_bin + 1, :],\n",
    "                    axis=0,\n",
    "                    keepdims=True,\n",
    "                )\n",
    "                baseline_vec = np.nan_to_num(baseline_vec, nan=0.0, posinf=0.0, neginf=0.0)\n",
    "                curr_sp_trial = curr_sp_trial - baseline_vec\n",
    "                all_trials_data[:, :, itr] = curr_sp_trial.T\n",
    "\n",
    "            data_for_neuron_pca.append(np.nanmean(all_trials_data, axis=2).T.astype(np.float32, copy=False))\n",
    "            concat_all_trials_data.append(all_trials_data)\n",
    "            cond_id.append(np.full((n_trials_cond,), icond + 1, dtype=np.int16))\n",
    "\n",
    "        if not data_for_neuron_pca or not concat_all_trials_data:\n",
    "            continue\n",
    "\n",
    "        data_for_neuron_pca_mat = np.concatenate(data_for_neuron_pca, axis=0).astype(np.float32, copy=False)\n",
    "        if data_for_neuron_pca_mat.size == 0 or np.nansum(np.abs(data_for_neuron_pca_mat)) == 0:\n",
    "            continue\n",
    "\n",
    "        concat_all_trials_data_mat = np.concatenate(concat_all_trials_data, axis=2).astype(np.float32, copy=False)\n",
    "        cond_id_vec = np.concatenate(cond_id, axis=0).astype(np.int16, copy=False)\n",
    "\n",
    "        n_neurons = int(concat_all_trials_data_mat.shape[0])\n",
    "        n_all_trials = int(concat_all_trials_data_mat.shape[2])\n",
    "\n",
    "        pca = PCA(svd_solver='full')\n",
    "        pca.fit(data_for_neuron_pca_mat.astype(np.float64, copy=False))\n",
    "\n",
    "        coeff_neuron = pca.components_.T.astype(np.float32, copy=False)\n",
    "        explained_neuron = (pca.explained_variance_ratio_ * 100.0).astype(np.float32, copy=False)\n",
    "        mu_neuron = pca.mean_.astype(np.float32, copy=False)\n",
    "\n",
    "        if coeff_neuron.shape[1] < 2:\n",
    "            continue\n",
    "\n",
    "        print(\n",
    "            f'  session={session_id} probe={probe_id} -> '\n",
    "            f'PC1: {explained_neuron[0]:.1f}% | PC2: {explained_neuron[1]:.1f}%'\n",
    "        )\n",
    "\n",
    "        # Project every single trial onto the first two neuron-space PCs.\n",
    "        trial_trajectories_pc = np.zeros((n_all_trials, 2, n_timebins), dtype=np.float32)\n",
    "        for itrial in range(n_all_trials):\n",
    "            trial_data = concat_all_trials_data_mat[:, :, itrial].T.astype(np.float32, copy=False)\n",
    "            projected_trial = (trial_data - mu_neuron[np.newaxis, :]) @ coeff_neuron[:, :2]\n",
    "            trial_trajectories_pc[itrial, 0, :] = projected_trial[:, 0]\n",
    "            trial_trajectories_pc[itrial, 1, :] = projected_trial[:, 1]\n",
    "\n",
    "        session_record: Dict[str, Any] = {'conditions': []}\n",
    "        for i_condition in range(1, len(params['TrialType']) + 1):\n",
    "            curr_trial_id = np.flatnonzero(cond_id_vec == i_condition)\n",
    "            curr_trial_trajectories_pc = trial_trajectories_pc[curr_trial_id, :, :]\n",
    "\n",
    "            session_record['conditions'].append({\n",
    "                'probe_id': probe_id,\n",
    "                'condition': int(i_condition),\n",
    "                'n_trials': int(curr_trial_id.size),\n",
    "                'n_neurons': int(n_neurons),\n",
    "                'trial_trajectories_pc': curr_trial_trajectories_pc.astype(np.float32, copy=False),\n",
    "                'explained_variance': explained_neuron[:2].astype(np.float32, copy=False),\n",
    "                'sessionID': session_id,\n",
    "            })\n",
    "\n",
    "        sessions.append(session_record)\n",
    "\n",
    "    n_sessions = len(sessions)\n",
    "    area_avg = {'conditions': []}\n",
    "    for i_condition in range(1, len(params['TrialType']) + 1):\n",
    "        area_avg['conditions'].append({'n_sessions': int(n_sessions)})\n",
    "\n",
    "    attractor_results[current_area] = {\n",
    "        'sessions': sessions,\n",
    "        'area_avg': area_avg,\n",
    "    }\n",
    "\n",
    "    print(f'  -> Stored {n_sessions} valid sessions for {current_area}')\n",
    "\n",
    "print('Finished PC_projections export.')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Save\n",
    "\n",
    "Write the PCA projection payload to `processed_data/PC_projections.npz`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.savez_compressed(\n",
    "    out_path,\n",
    "    attractor_results=np.array([attractor_results], dtype=object),\n",
    "    params=np.array([params], dtype=object),\n",
    "    pca_endbin=np.int32(pca_endbin),\n",
    "    time_axis=time_axis.astype(np.float32, copy=False),\n",
    "    time_axis_full=time_axis_full.astype(np.float32, copy=False),\n",
    ")\n",
    "\n",
    "print(f'Saved PC-projection 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
}