{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# A gentle introduction to temporal response functions (TRFs)\n",
    "\n",
    "This tutorial builds a TRF analysis from scratch using simulated data. The goal is to make the shape of every object explicit:\n",
    "\n",
    "1. create stimulus features and an EEG-like response;\n",
    "2. fit and inspect a TRF;\n",
    "3. choose regularisation with cross-validation;\n",
    "4. use banded ridge when features have different scales or reliability; and\n",
    "5. encourage smooth temporal filters with quadratic regularisation.\n",
    "\n",
    "All figures are made in notebook cells, so they are included automatically in the HTML documentation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from pyeeg import TRFEstimator\n",
    "from pyeeg.simulate import (\n",
    "    dummy_trf_kernel,\n",
    "    simulate_pulse_inputs,\n",
    "    simulate_smooth_input,\n",
    "    simulate_trf_output,\n",
    ")\n",
    "\n",
    "plt.rcParams['figure.figsize'] = (10, 3)\n",
    "plt.rcParams['figure.dpi'] = 110"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Simulate two stimulus features\n",
    "\n",
    "A TRF maps a time-varying feature matrix `X` to an EEG matrix `y`. Here the first feature is a smooth, envelope-like signal and the second is a sparse word-onset-like signal. We use one known kernel for each feature, then add noise."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rng = np.random.default_rng(7)\n",
    "fs, duration = 50, 24\n",
    "t, envelope = simulate_smooth_input(duration, fs, fmax=5, seed=7)\n",
    "_, words = simulate_pulse_inputs(n_events=60, dur=duration, srate=fs, seed=8)\n",
    "X = np.column_stack([envelope, words])\n",
    "\n",
    "t_kernel, kernel = dummy_trf_kernel(-0.2, 0.6, fs, tloc=0.18, sigma=0.09)\n",
    "y = (\n",
    "    simulate_trf_output(t_kernel, 1.0 * kernel, envelope, fs)\n",
    "    + simulate_trf_output(t_kernel, 0.7 * kernel, words, fs)\n",
    "    + 0.25 * rng.standard_normal(len(t))\n",
    ")[:, None]\n",
    "\n",
    "fig, axes = plt.subplots(2, 1, sharex=True, figsize=(10, 4))\n",
    "axes[0].plot(t, envelope, label='envelope-like feature')\n",
    "axes[0].plot(t, words, label='word-onset feature', alpha=.8)\n",
    "axes[0].legend(loc='upper right')\n",
    "axes[0].set_ylabel('feature value')\n",
    "axes[1].plot(t, y[:, 0], color='black')\n",
    "axes[1].set(xlabel='time (s)', ylabel='EEG (a.u.)', title='Simulated one-channel EEG')\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Fit the simplest TRF\n",
    "\n",
    "`tmin` and `tmax` define the lags (in seconds) at which the response is estimated. `coef_` has shape `(n_lags, n_features, n_channels)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trf = TRFEstimator(tmin=-0.2, tmax=0.6, srate=fs, alpha=10, verbose=False)\n",
    "trf.fit(X, y, feat_names=['envelope', 'word onset'])\n",
    "print('X:', X.shape, 'y:', y.shape, 'coef_:', trf.coef_.shape)\n",
    "print(f'train correlation: {trf.score(X, y, reduce_multi=\"mean\"):.2f}')\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "for feature, name in enumerate(trf.feat_names_):\n",
    "    ax.plot(trf.times, trf.coef_[:, feature, 0], label=name)\n",
    "ax.axvline(0, color='black', lw=.8)\n",
    "ax.set(xlabel='lag (s)', ylabel='TRF weight', title='Estimated temporal response functions')\n",
    "ax.legend()\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Regularisation and cross-validation\n",
    "\n",
    "A larger ridge `alpha` shrinks coefficients and is useful when the lagged design matrix is noisy or collinear. Pass several values and call `xfit` to select the value with the best held-out correlation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "alphas = np.logspace(-2, 3, 7)\n",
    "cv_trf = TRFEstimator(tmin=-0.2, tmax=0.6, srate=fs, alpha=alphas, verbose=False)\n",
    "scores, best_alpha = cv_trf.xfit(X, y, n_splits=4, plot=False)\n",
    "mean_scores = scores.mean(axis=(0, 1, 3))\n",
    "print('best alpha:', best_alpha)\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.semilogx(alphas, mean_scores, 'o-')\n",
    "ax.axvline(best_alpha, color='tab:red', ls='--', label=f'best = {best_alpha:g}')\n",
    "ax.set(xlabel='ridge alpha', ylabel='cross-validated correlation', title='Choosing regularisation')\n",
    "ax.legend()\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Banded ridge: one penalty per feature\n",
    "\n",
    "Envelope and word-level features need not deserve the same penalty. `feature_alphas` supplies one ridge strength per input column. The next fit deliberately penalises the sparse word feature more strongly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "banded = TRFEstimator(\n",
    "    tmin=-0.2, tmax=0.6, srate=fs,\n",
    "    feature_alphas=[1, 100], verbose=False,\n",
    ")\n",
    "banded.fit(X, y, feat_names=['envelope', 'word onset'])\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "for feature, name in enumerate(banded.feat_names_):\n",
    "    ax.plot(banded.times, banded.coef_[:, feature, 0], label=f'{name} (alpha={banded.feature_alphas[feature]:g})')\n",
    "ax.axvline(0, color='black', lw=.8)\n",
    "ax.set(xlabel='lag (s)', ylabel='TRF weight', title='Banded-ridge TRF')\n",
    "ax.legend()\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Smoothness regularisation\n",
    "\n",
    "Quadratic regularisation penalises changes between neighbouring lags rather than simply shrinking every coefficient. Compare an ordinary ridge fit with `quadratic_reg='smoothness'`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ridge = TRFEstimator(tmin=-0.2, tmax=0.6, srate=fs, alpha=30, verbose=False)\n",
    "smooth = TRFEstimator(tmin=-0.2, tmax=0.6, srate=fs, alpha=30, quadratic_reg='smoothness', verbose=False)\n",
    "ridge.fit(X, y)\n",
    "smooth.fit(X, y)\n",
    "\n",
    "fig, axes = plt.subplots(1, 2, sharey=True, figsize=(10, 3))\n",
    "for feature, name in enumerate(['envelope', 'word onset']):\n",
    "    axes[feature].plot(ridge.times, ridge.coef_[:, feature, 0], label='ridge')\n",
    "    axes[feature].plot(smooth.times, smooth.coef_[:, feature, 0], label='quadratic smoothness')\n",
    "    axes[feature].axvline(0, color='black', lw=.8)\n",
    "    axes[feature].set(title=name, xlabel='lag (s)')\n",
    "axes[0].set_ylabel('TRF weight')\n",
    "axes[1].legend()\n",
    "fig.suptitle('Two kinds of regularisation')\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Timing different fitting paths\n",
    "\n",
    "The SVD ridge path is usually the quickest choice for ordinary fits. Quadratic regularisation adds a structured matrix solve, while robust Cauchy fitting iterates several weighted fits. This small benchmark is only illustrative: timings depend on the machine and data size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def elapsed(estimator):\n",
    "    start = time.perf_counter()\n",
    "    estimator.fit(X, y)\n",
    "    return time.perf_counter() - start\n",
    "\n",
    "timed = {\n",
    "    'ridge (SVD)': TRFEstimator(tmin=-.2, tmax=.6, srate=fs, alpha=30, verbose=False),\n",
    "    'smoothness': TRFEstimator(tmin=-.2, tmax=.6, srate=fs, alpha=30, quadratic_reg='smoothness', verbose=False),\n",
    "    'robust IRLS': TRFEstimator(tmin=-.2, tmax=.6, srate=fs, loss='cauchy', robust_sigma=.25, robust_max_iter=5, verbose=False),\n",
    "}\n",
    "times = {name: elapsed(estimator) for name, estimator in timed.items()}\n",
    "for name, seconds in times.items():\n",
    "    print(f'{name:16s}: {seconds:.3f} s')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Take-home messages\n",
    "\n",
    "* Keep `X` as `(samples, features)` and `y` as `(samples, channels)`.\n",
    "* Use `alpha` (and cross-validation) for ordinary ridge shrinkage.\n",
    "* Use `feature_alphas` for banded ridge when feature groups have different reliability.\n",
    "* Use `quadratic_reg='smoothness'` when neighbouring TRF lags should vary smoothly.\n",
    "* Inspect the estimated filters and held-out scores, not only the training fit."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"},
  "language_info": {"name": "python", "version": "3.10"}
 },
 "nbformat": 4,
 "nbformat_minor": 5
}