{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b457b7e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import h5py\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "from matplotlib.ticker import MaxNLocator\n",
    "import warnings\n",
    "import vbn_utils\n",
    "import decoding_utils as du\n",
    "import ccf_utils\n",
    "from vbn_utils import formatFigure\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1de7afae",
   "metadata": {},
   "outputs": [],
   "source": [
    "from notebook_utils import (\n",
    "    calc_dprime,\n",
    "    time_to_threshold_from_baseline_back_from_peak,\n",
    "    conf_interval\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1bb312e9",
   "metadata": {},
   "source": [
    "## Data loading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c717608",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Paths to all of the useful supplemental tables and tensors\n",
    "active_tensor_file = \"/Volumes/programs/mindscope/workgroups/np-behavior/vbn_data_release/supplemental_tables/vbnAllUnitSpikeTensor.hdf5\"\n",
    "\n",
    "stim_table_file = \"/Volumes/programs/mindscope/workgroups/np-behavior/vbn_data_release/supplemental_tables/master_stim_table_no_filter.csv\"\n",
    "unit_table_file = \"/Volumes/programs/mindscope/workgroups/np-behavior/vbn_data_release/supplemental_tables/master_units_with_responsiveness.csv\"\n",
    "\n",
    "sessions_table_file = \"/Volumes/programs/mindscope/workgroups/np-behavior/vbn_data_release/supplemental_tables/master_sessions_table.csv\" #\"/Volumes/programs/mindscope/workgroups/np-behavior/vbn_data_release/vbn_s3_cache/visual-behavior-neuropixels-0.5.0/project_metadata/ecephys_sessions.csv\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "167bb232",
   "metadata": {},
   "outputs": [],
   "source": [
    "units = pd.read_csv(unit_table_file)\n",
    "units['cortical_layer'] = units['cortical_layer'].replace('3-Feb','2/3') # 2/3 sometimes gets incorrectly reformatted as a date\n",
    "\n",
    "stim_table = pd.read_csv(stim_table_file)\n",
    "stim_table = stim_table.drop(columns='Unnamed: 0')\n",
    "\n",
    "active_tensor = h5py.File(active_tensor_file)\n",
    "\n",
    "sessions_table = pd.read_csv(sessions_table_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5c26289",
   "metadata": {},
   "outputs": [],
   "source": [
    "structure_tree = pd.read_csv(\"/Volumes/programs/mindscope/workgroups/np-behavior/vbn_data_release/supplemental_tables/ccf_structure_tree_2017.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18648bcf",
   "metadata": {},
   "outputs": [],
   "source": [
    "g_images = ['omitted'] + list(np.sort(stim_table[(stim_table['stimulus_name'].str.contains('_G_'))&\n",
    "                    (~stim_table['omitted'])&\n",
    "                    (~stim_table['image_name'].isin(['im083_r','im111_r']))]['image_name'].unique())) + ['im083_r','im111_r']\n",
    "\n",
    "h_images = ['omitted'] + list(np.sort(stim_table[(stim_table['stimulus_name'].str.contains('_H_'))&\n",
    "                    (~stim_table['omitted'])&\n",
    "                    (~stim_table['image_name'].isin(['im083_r','im111_r']))]['image_name'].unique())) + ['im083_r','im111_r']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa5db3d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_dict = {'G': g_images, 'H': h_images}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de04a46c",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.rcParams['font.size'] = 14\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce4061d5",
   "metadata": {},
   "source": [
    "## Calculate proportion of cells responsive over time for each area"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f40470e6",
   "metadata": {},
   "source": [
    "### Compute (if pre-computed, load below)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e08826c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "unit_ids = vbn_utils.get_unit_ids(units, ['VISall', 'SCMRN', 'LGd', 'LP'])\n",
    "session_list = units.set_index('unit_id').loc[unit_ids]['ecephys_session_id'].unique()\n",
    "\n",
    "session_responsiveness_over_time = vbn_utils.unit_responsiveness_over_time(active_tensor_file, session_list, unit_ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8b6a2e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "rot_dict = {}\n",
    "for srot in session_responsiveness_over_time:\n",
    "    unitid = srot['unit_ids'][0]\n",
    "    session_id = units.set_index('unit_id').loc[unitid]['ecephys_session_id']\n",
    "    rot_dict[session_id] = srot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fe47b10",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "\n",
    "with open(\"/Volumes/programs/mindscope/workgroups/np-behavior/vbn_data_release/supplemental_tables/responsiveness_over_time_dict.pkl\", 'wb') as file:\n",
    "    pickle.dump(rot_dict, file)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b02ba02",
   "metadata": {},
   "source": [
    "### Load if pre-computed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11f3bd7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "\n",
    "with open(\"/Volumes/programs/mindscope/workgroups/np-behavior/vbn_data_release/supplemental_tables/responsiveness_over_time_dict.pkl\", 'rb') as file:\n",
    "    rot_dict = pickle.load(file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a1b2145",
   "metadata": {},
   "outputs": [],
   "source": [
    "no_abnorm = sessions_table['abnormal_activity'].isnull() & sessions_table['abnormal_histology'].isnull()\n",
    "novel_sessions = sessions_table[(sessions_table['experience_level']=='Novel')&no_abnorm]['ecephys_session_id'].values\n",
    "familiar_sessions = sessions_table[(sessions_table['experience_level']=='Familiar')&no_abnorm]['ecephys_session_id'].values"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37bf07b7",
   "metadata": {},
   "source": [
    "### Plot responsiveness over time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3a54264",
   "metadata": {},
   "outputs": [],
   "source": [
    "region_names = ['LGd', 'VISp', 'VISl', 'VISrl', 'VISal', 'LP', 'VISpm', 'VISam', 'Hipp', 'SCMRN', 'VISall']\n",
    "\n",
    "cell_type = 'all'\n",
    "plt.rcParams['font.size'] = 16\n",
    "region_summary_rot = {r:{'Novel_unshared':[], 'Novel_shared':[], 'Familiar_unshared':[], 'Familiar_shared':[]} for r in region_names}\n",
    "for region in region_names:\n",
    "    for sessions, condition in zip([novel_sessions, familiar_sessions], ['Novel', 'Familiar']):\n",
    "        for session_id in sessions:\n",
    "            \n",
    "            session_rot = rot_dict[session_id]\n",
    "            \n",
    "            units_to_analyze = vbn_utils.get_unit_ids(units, region, cell_types=cell_type, session_id=session_id, \n",
    "                                                            training_trajectory='all')\n",
    "            units_to_analyze = units[units['unit_id'].isin(units_to_analyze)]\n",
    "            \n",
    "            if len(units_to_analyze)==0:\n",
    "                continue\n",
    "            \n",
    "            response_inds = np.isin(session_rot['unit_ids'], units_to_analyze['unit_id'].values)\n",
    "\n",
    "            shared_images = ['im083_r', 'im111_r']\n",
    "            non_shared_images = [k for k,v in session_rot.items() if 'im' in k and ~np.isin(k, shared_images)]\n",
    "\n",
    "            for image_type_set, image_label in zip([shared_images, non_shared_images], ['_shared', '_unshared']):\n",
    "                image_type_data = []\n",
    "                for im in image_type_set:\n",
    "                    im_sparseness = session_rot[im][response_inds]\n",
    "                    sig = im_sparseness<0.01\n",
    "                    image_type_data.append(np.mean(sig, axis=0))\n",
    "                \n",
    "                region_summary_rot[region][f'{condition}{image_label}'].append(np.mean(image_type_data, axis=0))\n",
    "\n",
    "\n",
    "    fig, axes = plt.subplots(1,2)\n",
    "    fig.set_size_inches([8,4.5])\n",
    "    ax = axes[0]\n",
    "    time = np.arange(100) + 20\n",
    "    color_dict = {'Familiar_unshared': 'b', 'Novel_shared': 'purple', 'Novel_unshared': 'r'}\n",
    "    for key in ['Familiar_unshared','Novel_unshared']:\n",
    "        num_no_nan = np.sum([np.all(~np.isnan(r)) for r in region_summary_rot[region][key]])\n",
    "\n",
    "        mean = np.nanmean(region_summary_rot[region][key], axis=0)[:100]\n",
    "        sem = np.nanstd(region_summary_rot[region][key], axis=0)/num_no_nan**0.5\n",
    "        sem = sem[:100]\n",
    "        ax.plot(time, mean, color=color_dict[key], label=key)\n",
    "        ax.fill_between(time, mean+sem, mean-sem, color=color_dict[key], alpha=0.5)\n",
    "\n",
    "        vals = np.stack(region_summary_rot[region][key])\n",
    "        for itimepoint, timepoint in enumerate([20, 60]):\n",
    "            time_mean = np.nanmean(vals[:, timepoint])\n",
    "            time_sem = np.nanstd(vals[:, timepoint])/num_no_nan**0.5\n",
    "            axes[1].plot(itimepoint, time_mean, color=color_dict[key], marker='o', markersize=10)\n",
    "            axes[1].errorbar(itimepoint, time_mean, yerr=time_sem, color=color_dict[key], capsize=5)\n",
    "    axes[1].set_xticks([0, 1])\n",
    "    axes[1].set_xticklabels(['20-60ms', '60-100ms'])   \n",
    "    \n",
    "    ax.set_title(region)\n",
    "    ax.set_xlabel('Time from stim start (ms)')\n",
    "    ax.set_ylabel('Fraction cells responsive')\n",
    "    ymax = ax.get_ylim()[1]\n",
    "    [ax.plot(xval, ymax, 'kv', mfc=mfc, ms=10) for mfc, xval in zip(['k', 'w'], [40, 80])]\n",
    "    \n",
    "    ax.set_ylim([0, ax.get_ylim()[1]])\n",
    "    ax.set_xticks((20, 100))\n",
    "    ax.spines['bottom'].set_bounds(20, 100)\n",
    "    for spine in ax.spines.values():\n",
    "        spine.set_linewidth(1.5)  # Set spine linewidth to 1.5 points\n",
    "\n",
    "    # Set the linewidth of ticks\n",
    "    ax.tick_params(width=1.5) \n",
    "    [vbn_utils.formatFigure(fig, a) for a in axes]\n",
    "    ax.yaxis.set_major_locator(MaxNLocator(nbins=5)) # 5 ticks on the y-axis\n",
    "    ax.xaxis.set_label_coords(0.4, -0.12)\n",
    "    plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db02af45",
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib.lines import Line2D\n",
    "\n",
    "plt.rcParams['font.size'] = 18\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "regions = ['LGd', 'LP', 'VISp', 'VISl', 'VISrl', 'VISal', 'VISpm', 'VISam']\n",
    "for ir, region in enumerate(regions):\n",
    "    \n",
    "    iteration_timepoint_vals = np.full((2, 2, 1000), np.nan)\n",
    "    for ikey, key in enumerate(['Familiar_unshared', 'Novel_unshared']):\n",
    "        vals = np.stack(region_summary_rot[region][key])\n",
    "\n",
    "        for iteration in range(1000):\n",
    "            iteration_indices = np.random.choice(np.arange(vals.shape[0]), vals.shape[0], replace=True)\n",
    "            for itimepoint, timepoint in enumerate([20, 60]):\n",
    "                time_mean = np.nanmean(vals[iteration_indices, timepoint])\n",
    "                iteration_timepoint_vals[ikey, itimepoint, iteration] = time_mean\n",
    "        \n",
    "    \n",
    "    iteration_diffs = iteration_timepoint_vals[1] - iteration_timepoint_vals[0]\n",
    "    iteration_means = iteration_diffs.mean(axis=1)\n",
    "\n",
    "    ci_high = np.percentile(iteration_diffs, 97.5, axis=1)\n",
    "    ci_low = np.percentile(iteration_diffs, 2.5, axis=1)\n",
    "\n",
    "    color = ccf_utils.get_area_color(region, structure_tree)\n",
    "    for mfc, ind in zip([color, 'w'], [0,1]):\n",
    "        ax.plot(ir, iteration_means[ind], 'o', color=color, mfc=mfc, ms=11, markeredgewidth=1.5)\n",
    "        ax.errorbar([ir], [iteration_means[ind]], yerr=[[ci_high[ind]-iteration_means[ind]], [iteration_means[ind]-ci_low[ind]]], \n",
    "                        color=color)\n",
    "\n",
    "ax.axhline(0, color='k', ls='dotted')\n",
    "ax.set_xticks(np.arange(len(regions)))\n",
    "ax.set_xticklabels(regions, rotation=90)\n",
    "ax.set_ylabel(r\"$\\Delta$ fraction responsive (Nov - Fam)\")\n",
    "for spine in ax.spines.values():\n",
    "    spine.set_linewidth(1.5)  # Set spine linewidth to 1.5 points\n",
    "\n",
    "# Set the linewidth of ticks\n",
    "ax.tick_params(width=1.5) \n",
    "\n",
    "custom_lines = [\n",
    "    Line2D([0], [0], marker='o', color='k', markerfacecolor='w', markersize=10, lw=0, label='late'),\n",
    "    Line2D([0], [0], marker='o', color='k', markerfacecolor='k', markersize=10, lw=0, label='early'),\n",
    "]\n",
    "\n",
    "plt.legend(handles=custom_lines, loc='upper left', frameon=False, bbox_to_anchor=(-0.05, 1.1), handletextpad=0)\n",
    "vbn_utils.formatFigure(fig, ax)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ceda3e1e",
   "metadata": {},
   "source": [
    "## Load precomputed PSTH data for each unit (created in Figure 4 notebook)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c23b005",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "\n",
    "with open(\"/Volumes/programs/mindscope/workgroups/np-behavior/vbn_data_release/supplemental_tables/change_prechange_responses_active_passive_by_unitid.pkl\", 'rb') as file:\n",
    "    flash_data = pickle.load(file)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8466d5dd",
   "metadata": {},
   "source": [
    "## Calculate visual latency and novelty modulation latency for each area"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a951881b",
   "metadata": {},
   "source": [
    "### Compute latencies (if pre-computed, load below)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0bd668f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from analysis_utils import exponential_convolve\n",
    "\n",
    "base_slice = slice(0, 50)\n",
    "base_sub = lambda x: x - x[:, base_slice].mean(axis=1)[:, None]\n",
    "\n",
    "regions = ['LGd',\n",
    "    'LP',\n",
    "    'VISp',\n",
    "    'VISl',\n",
    "    'VISal',\n",
    "    'VISrl',\n",
    "    'VISpm',\n",
    "    'VISam',\n",
    "    'SCMRN',\n",
    "    'VISall'\n",
    "    ]\n",
    "\n",
    "clusters_to_plot = 'all' \n",
    "\n",
    "conditions = ['Familiar', 'Novel', 'Familiar_omitted', 'Novel_omitted']\n",
    "\n",
    "baseline_samples = 50\n",
    "\n",
    "latencies = {r:{c:{'Familiar': [], 'Novel': [], 'dprime': []} for c in ['all', 'RS', 'FS', 'SST']} for r in regions}\n",
    "for cell_type in ['all', 'RS', 'FS', 'SST']:\n",
    "    for region in regions:\n",
    "        if cell_type != 'all' and 'VIS' not in region:\n",
    "            continue\n",
    "\n",
    "        data_dict = {}\n",
    "        for condition in conditions:\n",
    "            if 'omitted' in condition:\n",
    "                image_category = 'omission'\n",
    "            else:\n",
    "                image_category = 'nonshared_nonchange'\n",
    "\n",
    "            units_to_use = vbn_utils.get_unit_ids(units, region, cell_types=cell_type, layers='all', clusters='all', clustering = 'new', \n",
    "                    experience=condition.split('_')[0], responsive=False, session_id='all')\n",
    "            \n",
    "            data = [flash_data[u]['active'][image_category] for u in units_to_use]\n",
    "            \n",
    "            data = np.array([exponential_convolve(d, 3, symmetrical=True) for d in data])\n",
    "            data_dict[condition] = base_sub(data)\n",
    "        \n",
    "        fam = data_dict['Familiar']\n",
    "        nov = data_dict['Novel']\n",
    "        fam_omitted = data_dict['Familiar_omitted']\n",
    "        nov_omitted = data_dict['Novel_omitted']\n",
    "        \n",
    "        num_fam = fam.shape[0]\n",
    "        num_nov = nov.shape[0]\n",
    "\n",
    "        dprime_iterations = []\n",
    "        dprime_null_iterations = []\n",
    "        diff_iterations = []\n",
    "        time_to_thresh = []\n",
    "        for iteration in range(1000):\n",
    "            dprime_over_time = []\n",
    "            dprime_fvso_over_time = []\n",
    "            dprime_nvso_over_time = []\n",
    "\n",
    "            sample_fam = fam[np.random.choice(np.arange(num_fam), num_fam, replace=True)]\n",
    "            sample_nov = nov[np.random.choice(np.arange(num_nov), num_nov, replace=True)]\n",
    "            sample_fam_omitted = fam_omitted[np.random.choice(np.arange(num_fam), num_fam, replace=True)]\n",
    "            sample_nov_omitted = nov_omitted[np.random.choice(np.arange(num_nov), num_nov, replace=True)]\n",
    "            \n",
    "\n",
    "            for timepoint in range(fam.shape[1]):\n",
    "                dp = calc_dprime(sample_nov[:, timepoint], sample_fam[:, timepoint], signed=True)\n",
    "                dprime_over_time.append(dp)\n",
    "\n",
    "                dp_fvso = calc_dprime(sample_fam[:, timepoint], sample_fam_omitted[:, timepoint], signed=True)\n",
    "                dp_nvso = calc_dprime(sample_nov[:, timepoint], sample_nov_omitted[:, timepoint], signed=True)\n",
    "                dprime_fvso_over_time.append(dp_fvso)\n",
    "                dprime_nvso_over_time.append(dp_nvso)\n",
    "\n",
    "\n",
    "            latencies[region][cell_type]['dprime'].append(time_to_threshold_from_baseline_back_from_peak(dprime_over_time[baseline_samples:], threshold=0.1)[0] + baseline_samples)\n",
    "            latencies[region][cell_type]['Familiar'].append(time_to_threshold_from_baseline_back_from_peak(dprime_fvso_over_time[baseline_samples:], threshold=0.1)[0] + baseline_samples)\n",
    "            latencies[region][cell_type]['Novel'].append(time_to_threshold_from_baseline_back_from_peak(dprime_nvso_over_time[baseline_samples:], threshold=0.1)[0] + baseline_samples)\n",
    "\n",
    "            dprime_iterations.append(dprime_over_time)\n",
    "\n",
    "        median_dprime_over_samples = np.median(dprime_iterations, axis=0)\n",
    "        cis = np.array([conf_interval(vals) for vals in np.array(dprime_iterations).T])\n",
    "\n",
    "        fam_mean = np.nanmean(fam, axis=0)\n",
    "        nov_mean = np.nanmean(nov, axis=0)\n",
    "\n",
    "        median_dprime_time_to_thresh = np.nanmedian(latencies[region][cell_type]['dprime'])\n",
    "        fig, ax = plt.subplots()\n",
    "        fig.suptitle(f'{region} {cell_type} \\n fam n: {num_fam}, nov n: {num_nov}')\n",
    "        ax2 = ax.twinx()\n",
    "        ax.plot(median_dprime_over_samples, 'k')\n",
    "        ax.plot(median_dprime_time_to_thresh, median_dprime_over_samples[int(median_dprime_time_to_thresh)], 'ko')\n",
    "        ax.fill_between(np.arange(len(median_dprime_over_samples)), cis[:,0], cis[:,1], color='k', alpha=0.3, lw=0)\n",
    "        ax.axvline(np.nanmedian(latencies[region][cell_type]['dprime']), color='k', ls='dotted')\n",
    "\n",
    "        ax2.plot(np.nanmean(fam, axis=0), 'b')\n",
    "        ax2.plot(np.nanmean(nov, axis=0), 'r')\n",
    "        ax2.plot(np.nanmean(fam_omitted, axis=0), 'b', alpha=0.5)\n",
    "        ax2.plot(np.nanmean(nov_omitted, axis=0), 'r', alpha=0.5)\n",
    "        ax2.axvline(np.nanmedian(latencies[region][cell_type]['Familiar']), color='b', ls='dotted')\n",
    "        ax2.axvline(np.nanmedian(latencies[region][cell_type]['Novel']), color='r', ls='dotted')\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5128c60e",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"/Volumes/programs/mindscope/workgroups/np-behavior/vbn_data_release/supplemental_tables/latencies_by_region_and_celltype.pkl\", 'wb') as file:\n",
    "    pickle.dump(latencies, file)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d0bff80",
   "metadata": {},
   "source": [
    "### Load pre-computed latencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce49e9a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load latencies for plotting\n",
    "with open(\"/Volumes/programs/mindscope/workgroups/np-behavior/vbn_data_release/supplemental_tables/latencies_by_region_and_celltype.pkl\", 'rb') as file:\n",
    "    latencies = pickle.load(file)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efd28389",
   "metadata": {},
   "source": [
    "### Plot visual and novelty modulation latencies vs. anatomical hierarchy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17a6e207",
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy.stats\n",
    "\n",
    "plt.rcParams.update({'font.size':14.5})\n",
    "\n",
    "hierarchy_dict ={\n",
    "                    'LGd':\t-0.515027963,\n",
    "                    'VISp':\t-0.357332099,\n",
    "                    'VISl':\t-0.093888551,\n",
    "                    'VISrl':-0.059871325,\n",
    "                    'LP':\t0.10524781,\n",
    "                    'VISal':0.152217979,\n",
    "                    'VISpm':0.327668075,\n",
    "                    'VISam':0.440986074\n",
    "                }\n",
    "\n",
    "fig, axes = plt.subplots(2,1)\n",
    "fig.set_size_inches(5,10)\n",
    "ax = axes[0]\n",
    "ax2 = axes[1]\n",
    "\n",
    "binsize = 1\n",
    "vis_latencies = []\n",
    "dprime_latencies = []\n",
    "hierarchy_positions = []\n",
    "for ir, region in enumerate(['LGd', 'VISp', 'VISl', 'VISrl', 'VISal', 'LP', 'VISpm', 'VISam']):\n",
    "        visual_latency_array = np.stack([latencies[region]['all'][cond] for cond in ['Familiar', 'Novel']])\n",
    "        visual_latency_mean_over_cond = np.mean(visual_latency_array, axis = 0) - 50\n",
    "        visual_latency_mean_over_cond[(visual_latency_mean_over_cond>270)|(visual_latency_mean_over_cond<20)] = np.nan\n",
    "        vis_latencies.append(visual_latency_mean_over_cond)\n",
    "        dprime_latency = np.array(latencies[region]['all']['dprime']) - 50\n",
    "        dprime_latency[(dprime_latency>270)|(dprime_latency<20)] = np.nan\n",
    "        dprime_latencies.append(dprime_latency)\n",
    "        hierarchy_positions.append(hierarchy_dict[region])\n",
    "        ax.plot(hierarchy_dict[region], np.nanmean(visual_latency_mean_over_cond), 'o', color=ccf_utils.get_area_color(region, structure_tree), ms=10)\n",
    "        ax.errorbar(hierarchy_dict[region], np.nanmean(visual_latency_mean_over_cond), np.nanstd(visual_latency_mean_over_cond), color=ccf_utils.get_area_color(region, structure_tree), ls='none')\n",
    "        \n",
    "        ax2.plot(hierarchy_dict[region], np.nanmean(dprime_latency), 'o', color=ccf_utils.get_area_color(region, structure_tree), ms=10)\n",
    "        ax2.errorbar(hierarchy_dict[region], np.nanmean(dprime_latency), np.nanstd(dprime_latency), color=ccf_utils.get_area_color(region, structure_tree), ls='none',)\n",
    "\n",
    "ax.set_ylabel('Visual Response Latency (ms)')\n",
    "ax2.set_ylabel('Novelty Modulation Latency (ms)')\n",
    "\n",
    "xlims = np.array(ax.get_xlim())\n",
    "corr_params = scipy.stats.linregress(hierarchy_positions, np.nanmean(vis_latencies, axis=1))\n",
    "ax.plot(xlims, xlims*corr_params[0] + corr_params[1], 'k--',)\n",
    "\n",
    "ax.set_xticks([val for key, val in hierarchy_dict.items()])\n",
    "ax.set_xticklabels([f'{np.round(val, 2)} \\n{key}' if key not in ['VISrl', 'VISal', 'VISam'] else f'\\n\\n{key}' for key, val in hierarchy_dict.items()])\n",
    "\n",
    "[a.set_title('') for a in [ax, ax2]]\n",
    "for corr_func in [scipy.stats.linregress, scipy.stats.spearmanr]:\n",
    "        vis_response_corr = corr_func(hierarchy_positions, np.nanmean(vis_latencies, axis=1))\n",
    "        dprime_corr = corr_func(hierarchy_positions, np.nanmean(dprime_latencies, axis=1))\n",
    "\n",
    "        print(f'Vis latency {corr_func.__name__}: {vis_response_corr}')\n",
    "        print(f'Dprime latency {corr_func.__name__}: {dprime_corr}')\n",
    "        atitle = ax.get_title()\n",
    "        ax.set_title(atitle + '\\n' + corr_func.__name__ + ' pval:' + str(vis_response_corr.pvalue))\n",
    "\n",
    "        atitle = ax2.get_title()\n",
    "\n",
    "ax2.set_xticks([val for key, val in hierarchy_dict.items()])\n",
    "ax2.set_xticklabels([f'{np.round(val, 2)} \\n{key}' if key not in ['VISrl', 'VISal', 'VISam'] else f'\\n\\n{key}' for key, val in hierarchy_dict.items()])\n",
    "\n",
    "[formatFigure(fig, a) for a in axes]\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e57d608",
   "metadata": {},
   "source": [
    "## Novelty modulation PSTHs for LGd and VIS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25426d6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "areas_ordered = ['LGd', 'VISall']\n",
    "\n",
    "\n",
    "for cell_type in ['RS', 'FS', 'SST', 'all']:\n",
    "    fig = vbn_utils.make_nov_mod_psth_figure(areas_ordered, layers='all', cell_types=cell_type, clusters='all', units_subsample=units, flash_data=flash_data, norm=False, state=['active',],\n",
    "                        flashes=['nonshared_nonchange', 'change'], training_trajectory='all', plot_pvalues=False, smoothing_kernel_width=3, return_fig=True)\n",
    "    axes = fig.get_axes()\n",
    "\n",
    "    for area, ax in zip(areas_ordered, [axes[0], axes[1]]):\n",
    "        if cell_type != 'all' and 'VIS' not in area:\n",
    "            continue\n",
    "        visual_latency_array = np.stack([latencies[area][cell_type][cond] for cond in ['Familiar', 'Novel']])\n",
    "        visual_latency_mean_over_cond = np.mean(visual_latency_array, axis = 0) - 50\n",
    "        visual_latency_mean_over_cond[(visual_latency_mean_over_cond>270)|(visual_latency_mean_over_cond<20)] = np.nan\n",
    "\n",
    "        dprime_latency = np.array(latencies[area][cell_type]['dprime']) - 50\n",
    "        dprime_latency[(dprime_latency>270)|(dprime_latency<20)] = np.nan\n",
    "\n",
    "        ylims = ax.get_ylim()\n",
    "        ax.plot(np.nanmean(visual_latency_mean_over_cond), ylims[1], 'kv', mfc='w', ms=10)\n",
    "        ax.plot(np.nanmean(dprime_latency), ylims[1], 'kv', ms=10)\n",
    "\n",
    "        ax.spines['bottom'].set_bounds(20, 100)\n",
    "        for spine in ax.spines.values():\n",
    "            spine.set_linewidth(1.5)  # Set spine linewidth to 1.5 points\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b188bedc",
   "metadata": {},
   "source": [
    "## Image-wise response metrics (tuning changes Familiar→Novel)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c21b958c",
   "metadata": {},
   "source": [
    "### Compute image-wise responses (if pre-computed, load from units table)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1af30d8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "unit_ids = vbn_utils.get_unit_ids(units, ['VISall', 'SCMRN', 'LGd', 'LP'])\n",
    "session_list = units.set_index('unit_id').loc[unit_ids]['ecephys_session_id'].unique()\n",
    "\n",
    "session_unit_imagewise_responses = vbn_utils.unit_imagewise_stats(active_tensor_file, session_list, unit_ids, baseline_slice=slice(700, 750), response_slice=slice(20,100))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc42fe32",
   "metadata": {},
   "outputs": [],
   "source": [
    "unit_imagewise_responses = {}\n",
    "for d in session_unit_imagewise_responses:\n",
    "    unit_imagewise_responses.update(d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "826b4726",
   "metadata": {},
   "outputs": [],
   "source": [
    "ggh = units.loc[du.apply_condition_filter(units, 'Familiar', 'GGH') | du.apply_condition_filter(units, 'Novel', 'GGH')]\n",
    "hhg = units.loc[du.apply_condition_filter(units, 'Familiar', 'HHG') | du.apply_condition_filter(units, 'Novel', 'HHG')]\n",
    "ghg = units.loc[du.apply_condition_filter(units, 'Familiar', 'GHG') | du.apply_condition_filter(units, 'Novel', 'GHG')]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aeef4b8e",
   "metadata": {},
   "source": [
    "### Tuning curve analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3abaa0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "vis_units = vbn_utils.get_unit_ids(ggh, 'VISall')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c1606b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_ids = h_images[:-2] + g_images[1:]\n",
    "plt.rcParams['font.size'] = 22\n",
    "\n",
    "#tuning curves\n",
    "colors = ['b', 'r']\n",
    "regions = ['LGd', 'LP', 'VISall', 'VISp', 'VISl', 'VISal', 'VISrl', 'VISpm', 'VISam']\n",
    "tuning_diff_dict = {}\n",
    "for cell_type in ['RS', 'FS']:\n",
    "    for region in regions:\n",
    "        \n",
    "        condition_curves = []\n",
    "        fig, ax = plt.subplots()\n",
    "        ax2 = ax.twinx()\n",
    "        for ic, condition in enumerate(['Familiar', 'Novel']):\n",
    "\n",
    "            units_to_use = vbn_utils.get_unit_ids(units, region, experience=condition, cell_types=cell_type)\n",
    "\n",
    "            tuning_curves = []\n",
    "            for u in units_to_use:\n",
    "                tuning = np.array([unit_imagewise_responses[u][im]['mean'] for im in image_ids[1:-2]])\n",
    "                tuning = tuning[~np.isnan(tuning)]\n",
    "                if len(tuning)<6:\n",
    "                    continue\n",
    "\n",
    "                tuning_curves.append(np.sort(tuning)*1000)\n",
    "\n",
    "            tuning_curves = np.array(tuning_curves)\n",
    "\n",
    "            iteration_curves = []\n",
    "            for iteration in range(1000):\n",
    "                indices = np.random.choice(np.arange(tuning_curves.shape[0]), tuning_curves.shape[0], replace=True)\n",
    "                iteration_curves.append(np.nanmean(tuning_curves[indices], axis=0))\n",
    "\n",
    "            condition_curves.append(iteration_curves)\n",
    "\n",
    "            vbn_utils.mean_sem_plot(tuning_curves, ax, x=np.arange(1, 7), color=colors[ic])\n",
    "        \n",
    "        ax.set_title(f'{region} {cell_type}')\n",
    "        \n",
    "        condition_curves = np.array(condition_curves)\n",
    "        \n",
    "        vbn_utils.mean_CI_plot(condition_curves[1]-condition_curves[0], ax2, x=np.arange(1,7), color='k', ls='dotted')\n",
    "        tuning_diff_dict[(region, cell_type)] = condition_curves[1] - condition_curves[0]\n",
    "        \n",
    "        ax.set_zorder(ax2.get_zorder()+1)\n",
    "        ax.patch.set_visible(False)\n",
    "\n",
    "        ax.set_xticks(np.arange(1, 7))\n",
    "        ax.set_xlabel('Image rank')\n",
    "        ax.set_ylabel('Firing rate (Hz)')\n",
    "        ax2.set_ylabel('Difference in firing rate (Hz)', rotation=270, labelpad=23)\n",
    "        \n",
    "        ax.yaxis.set_major_locator(MaxNLocator(nbins=5))  # 5 ticks on the y-axis\n",
    "        ax2.yaxis.set_major_locator(MaxNLocator(nbins=5))  # 5 ticks on the y-axis\n",
    "\n",
    "\n",
    "        for a in [ax, ax2]:\n",
    "            for spine in a.spines.values():\n",
    "                spine.set_linewidth(1.25)  # Set spine linewidth to 2 points\n",
    "\n",
    "            # Set the linewidth of ticks\n",
    "            a.tick_params(width=1.25)\n",
    "\n",
    "        [vbn_utils.formatFigure(fig, a, yaxis_side=side) for a,side in zip([ax, ax2], ['left', 'right'])]\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "vbn_manuscript",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.8.20"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}