{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a3c65a72",
   "metadata": {},
   "source": [
    "# 📦 Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 149,
   "id": "c8ba79b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import animation\n",
    "from scipy.io import loadmat\n",
    "from IPython.display import HTML\n",
    "import imageio"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b163b96e",
   "metadata": {},
   "source": [
    "# 📑 Load metadata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 150,
   "id": "35e5984f",
   "metadata": {},
   "outputs": [],
   "source": [
    "metadata_path = '../Metadata/OLST_Attempts.csv'\n",
    "attempts_df = pd.read_csv(metadata_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5cb051e6",
   "metadata": {},
   "source": [
    "# 🔢 USER INPUT: Enter OLST_attempt_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 151,
   "id": "109654f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "attempt_id = \"36_TRLGL_RR_V4_an2\"  # ENTER Example attempt ID\n",
    "participant_id = attempt_id.split('_')[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d1e0e86",
   "metadata": {},
   "source": [
    "# 📁 Define base directory paths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 152,
   "id": "c3b4f3ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "mocap_base = f'../Raw/mocap/{participant_id}'\n",
    "fp_base = f'../Raw/ForcePlate/{participant_id}'\n",
    "radar_base = f'../Processed/Radar_RDMs/{participant_id}'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f19397ce",
   "metadata": {},
   "source": [
    "# 🔍 Lookup metadata row"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 153,
   "id": "6a9ef760",
   "metadata": {},
   "outputs": [],
   "source": [
    "row = attempts_df[attempts_df['OLST_attempt_id'] == attempt_id].iloc[0]\n",
    "radar_id = row['RADAR_capture']\n",
    "an_number = row['an']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ede232ab",
   "metadata": {},
   "source": [
    "# 🧠 Inferred filenames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 154,
   "id": "0d2da7e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "mocap_file = f\"{radar_id.replace('_RR_', '_MC_')}_pos.csv\"\n",
    "foot_side = 'left' if any(tag in radar_id for tag in ['MNTRL', 'TRLG']) else 'right'\n",
    "fp_file = f\"{radar_id.replace('_RR_', '_FP_')}_{foot_side}.csv\"\n",
    "radar_file = f\"{radar_id}.mat\"\n",
    "\n",
    "mocap_path = os.path.join(mocap_base, mocap_file)\n",
    "fp_path = os.path.join(fp_base, fp_file)\n",
    "radar_path = os.path.join(radar_base, radar_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba158ad3",
   "metadata": {},
   "source": [
    "# 📈 Load sensor data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 155,
   "id": "4c82d5b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_mocap = pd.read_csv(mocap_path)\n",
    "df_fp = pd.read_csv(fp_path)\n",
    "mat_data = loadmat(radar_path)\n",
    "rdm_cube = mat_data['zoomed_RDM_cube']  # (range, doppler, frames)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15853735",
   "metadata": {},
   "source": [
    "# 🧭 Sync info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 156,
   "id": "8b190f58",
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- Time sync constants ---\n",
    "t_start = row['t_foot_up']\n",
    "t_stop = row['t_end']\n",
    "seconds_per_frame = row['Seconds_per_Frame']\n",
    "num_frames = int(np.floor((t_stop - t_start) / seconds_per_frame))\n",
    "\n",
    "# --- Create target time vector (radar-aligned) ---\n",
    "target_times = t_start + np.arange(num_frames) * seconds_per_frame\n",
    "\n",
    "# --- Resample MOCAP ---\n",
    "df_mocap['time'] = pd.to_numeric(df_mocap['time'], errors='coerce')\n",
    "df_mocap_interp = df_mocap.set_index('time').interpolate(method='linear', axis=0)\n",
    "df_mocap_sync = df_mocap_interp.reindex(target_times, method='nearest').reset_index(drop=True)\n",
    "\n",
    "# --- Resample FP ---\n",
    "df_fp['time'] = pd.to_numeric(df_fp['time'], errors='coerce')\n",
    "df_fp_interp = df_fp.set_index('time').interpolate(method='linear', axis=0)\n",
    "df_fp_sync = df_fp_interp.reindex(target_times, method='nearest').reset_index(drop=True)\n",
    "\n",
    "# --- Radar frame indices ---\n",
    "radar_frame_indices = list(range(int(row['frame_foot_up']), int(row['frame_foot_up']) + num_frames))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f4bf292",
   "metadata": {},
   "source": [
    "# ⚖️ Normalize force plate CoP for plotting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 157,
   "id": "65442c2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "cop_x = df_fp_sync['COP_X'].values\n",
    "cop_y = df_fp_sync['COP_Y'].values\n",
    "cop_x = (cop_x - np.min(cop_x)) / (np.max(cop_x) - np.min(cop_x))\n",
    "cop_y = (cop_y - np.min(cop_y)) / (np.max(cop_y) - np.min(cop_y))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d85f6955",
   "metadata": {},
   "source": [
    "# 🎥 Set up figure and animation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 158,
   "id": "15aee345",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 3, figsize=(15, 5))\n",
    "fig.tight_layout(rect=[0, 0, 1, 0.87])  # make room for a longer suptitle\n",
    "\n",
    "# Precompute time and frame ranges for the title\n",
    "time_start = target_times[0]\n",
    "time_end = target_times[-1]\n",
    "frame_start = radar_frame_indices[0]\n",
    "frame_end = radar_frame_indices[-1]\n",
    "\n",
    "def update(frame_idx):\n",
    "    f_radar = radar_frame_indices[frame_idx]\n",
    "    f_fp = frame_idx\n",
    "    f_mocap = frame_idx\n",
    "    current_time = target_times[frame_idx]\n",
    "\n",
    "    for ax in axes:\n",
    "        ax.clear()\n",
    "\n",
    "    # --- Radar RDM Frame ---\n",
    "    axes[0].imshow(rdm_cube[:, :, f_radar], aspect='auto', origin='lower', cmap='jet')\n",
    "    axes[0].set_title(f'Radar RDM\\nFrame {f_radar}')\n",
    "    axes[0].axis('off')\n",
    "\n",
    "    # --- Force Plate CoP ---\n",
    "    axes[1].plot(cop_x[:f_fp+1], cop_y[:f_fp+1], color='blue', alpha=0.7)\n",
    "    axes[1].scatter(cop_x[f_fp], cop_y[f_fp], color='red')\n",
    "    axes[1].set_xlim(0, 1)\n",
    "    axes[1].set_ylim(0, 1)\n",
    "    axes[1].set_title(f'Force Plate CoP\\nTime {current_time:.2f} s')\n",
    "    axes[1].axis('off')\n",
    "\n",
    "    # --- MOCAP Point Cloud Plot (Z-Y projection, excluding actuator markers) ---\n",
    "    pos_cols = [col for col in df_mocap_sync.columns if '_pos_' in col]\n",
    "    marker_names = sorted(set(col.rsplit('_pos_', 1)[0] for col in pos_cols))\n",
    "    marker_names = [name for name in marker_names if 'Actuator' not in name]\n",
    "\n",
    "    points = []\n",
    "    for marker in marker_names:\n",
    "        x_col = f\"{marker}_pos_X\"\n",
    "        y_col = f\"{marker}_pos_Y\"\n",
    "        z_col = f\"{marker}_pos_Z\"\n",
    "        if x_col in df_mocap_sync.columns and y_col in df_mocap_sync.columns and z_col in df_mocap_sync.columns:\n",
    "            x = df_mocap_sync.at[f_mocap, x_col]\n",
    "            y = df_mocap_sync.at[f_mocap, y_col]\n",
    "            z = df_mocap_sync.at[f_mocap, z_col]\n",
    "            points.append((x, y, z))\n",
    "\n",
    "    points = np.array(points)\n",
    "\n",
    "    if points.shape[0] > 0:\n",
    "        axes[2].scatter(points[:, 1], points[:, 2], color='purple', s=20)  # Z-Y projection\n",
    "        axes[2].set_title(f'MOCAP Point Cloud\\nZ-Y Projection\\nTime {current_time:.2f} s')\n",
    "        axes[2].axis('equal')\n",
    "        axes[2].axis('off')\n",
    "\n",
    "    # --- Super title with full info ---\n",
    "    fig.suptitle(\n",
    "        f\"{attempt_id} | Time {time_start:.2f}–{time_end:.2f} s | Frames {frame_start}–{frame_end}\",\n",
    "        fontsize=16\n",
    "    )\n",
    "\n",
    "    return axes\n",
    "\n",
    "ani = animation.FuncAnimation(fig, update, frames=num_frames, interval=1000 * seconds_per_frame)\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ec60598",
   "metadata": {},
   "source": [
    "# 💾 Save GIF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 159,
   "id": "48220f01",
   "metadata": {},
   "outputs": [],
   "source": [
    "gif_filename = f'viz_{attempt_id}.gif'\n",
    "ani.save(gif_filename, writer='pillow', dpi=100)\n",
    "plt.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "radartreepose_env",
   "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.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}