{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# PID Control Testing for Maglev Pod\n",
        "\n",
        "This notebook provides a feedforward + PID control interface for the `LevPodEnv` simulation environment.\n",
        "\n",
        "## Control Architecture\n",
        "- **Feedforward**: `MaglevPredictor` estimates the equilibrium PWM needed to hover at the current gap height\n",
        "- **Height PID**: Corrects residual gap error → additive PWM correction for all coils\n",
        "- **Roll PID**: Corrects roll angle → differential left/right adjustment\n",
        "- **Pitch PID**: Corrects pitch angle → differential front/back adjustment\n",
        "\n",
        "The outputs combine: `coil_pwm = feedforward + height_correction ± roll_adjustment ± pitch_adjustment`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 91,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Target Gap Height: 11.86 mm\n"
          ]
        }
      ],
      "source": [
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "from lev_pod_env import LevPodEnv, TARGET_GAP\n",
        "from pid_controller import PIDController, DEFAULT_GAINS\n",
        "from pid_simulation import run_pid_simulation, feedforward_pwm, build_feedforward_lut\n",
        "\n",
        "# Set plot style for better aesthetics\n",
        "plt.style.use('seaborn-v0_8-whitegrid')\n",
        "plt.rcParams['figure.facecolor'] = 'white'\n",
        "plt.rcParams['axes.facecolor'] = 'white'\n",
        "plt.rcParams['font.size'] = 11\n",
        "plt.rcParams['axes.titlesize'] = 14\n",
        "plt.rcParams['axes.labelsize'] = 12\n",
        "\n",
        "print(f\"Target Gap Height: {TARGET_GAP * 1000:.2f} mm\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "---\n",
        "## Load PID gains (JSON or defaults)\n",
        "\n",
        "Try to load `pid_best_params.json` (saved by Optuna in this folder). If missing, use `DEFAULT_GAINS`. You can override any value in the **PID Gains Configuration** cell below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 92,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Loaded PID gains from /Users/adipu/Documents/guadaloop_lev_control/RL Testing/pid_best_params.json\n"
          ]
        }
      ],
      "source": [
        "import os\n",
        "import json\n",
        "\n",
        "try:\n",
        "    import optuna_pid_tune as _optuna_tune\n",
        "    _best_path = os.path.join(os.path.dirname(os.path.abspath(_optuna_tune.__file__)), \"pid_best_params.json\")\n",
        "except Exception:\n",
        "    _best_path = \"pid_best_params.json\"\n",
        "\n",
        "if os.path.isfile(_best_path):\n",
        "    with open(_best_path) as f:\n",
        "        _best = json.load(f)\n",
        "    HEIGHT_KP = _best[\"height_kp\"]\n",
        "    HEIGHT_KI = _best[\"height_ki\"]\n",
        "    HEIGHT_KD = _best[\"height_kd\"]\n",
        "    ROLL_KP = _best[\"roll_kp\"]\n",
        "    ROLL_KI = _best[\"roll_ki\"]\n",
        "    ROLL_KD = _best[\"roll_kd\"]\n",
        "    PITCH_KP = _best[\"pitch_kp\"]\n",
        "    PITCH_KI = _best[\"pitch_ki\"]\n",
        "    PITCH_KD = _best[\"pitch_kd\"]\n",
        "    gains = dict(_best)\n",
        "    print(f\"Loaded PID gains from {_best_path}\")\n",
        "else:\n",
        "    HEIGHT_KP = DEFAULT_GAINS[\"height_kp\"]\n",
        "    HEIGHT_KI = DEFAULT_GAINS[\"height_ki\"]\n",
        "    HEIGHT_KD = DEFAULT_GAINS[\"height_kd\"]\n",
        "    ROLL_KP = DEFAULT_GAINS[\"roll_kp\"]\n",
        "    ROLL_KI = DEFAULT_GAINS[\"roll_ki\"]\n",
        "    ROLL_KD = DEFAULT_GAINS[\"roll_kd\"]\n",
        "    PITCH_KP = DEFAULT_GAINS[\"pitch_kp\"]\n",
        "    PITCH_KI = DEFAULT_GAINS[\"pitch_ki\"]\n",
        "    PITCH_KD = DEFAULT_GAINS[\"pitch_kd\"]\n",
        "    gains = dict(DEFAULT_GAINS)\n",
        "    print(\"Using DEFAULT_GAINS (no pid_best_params.json found). Run Optuna to save best params.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "---\n",
        "## PID Controller Class"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 93,
      "metadata": {},
      "outputs": [],
      "source": [
        "# PID controller and default gains live in pid_controller.py (imported above).\n",
        "# Override gains in the next cell; they are passed to run_pid_simulation via the gains dict."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "---\n",
        "## PID Gains Configuration (optional overrides)\n",
        "\n",
        "Constants are already set above from `pid_best_params.json` or defaults. To **override**, uncomment and edit a line in the cell below (e.g. `HEIGHT_KP = 50.0`) and re-run that cell; then run the simulation."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 94,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "PID Gains (used by simulation):\n",
            "  Height: Kp=100.96239556037482, Ki=0.0, Kd=8.173809263716658\n",
            "  Roll:   Kp=0.600856607986966, Ki=0.0, Kd=-0.1\n",
            "  Pitch:  Kp=50.011470879925824, Ki=0, Kd=1.8990306608320433\n"
          ]
        }
      ],
      "source": [
        "# Optional manual overrides: set any constant below and re-run this cell to use it.\n",
        "# Otherwise the values from the \"Load PID gains\" cell (JSON or DEFAULT_GAINS) are used.\n",
        "# HEIGHT_KP = 50.0\n",
        "# HEIGHT_KI = 5.0\n",
        "# HEIGHT_KD = 10.0\n",
        "# ROLL_KP = 2.0\n",
        "# ROLL_KI = 0.5\n",
        "# ROLL_KD = 0.5\n",
        "# PITCH_KP = 2.0\n",
        "# PITCH_KI = 0.5\n",
        "# PITCH_KD = 0.5\n",
        "\n",
        "# Build gains dict from current constants (used by run_pid_simulation)\n",
        "gains = {\n",
        "    \"height_kp\": HEIGHT_KP, \"height_ki\": HEIGHT_KI, \"height_kd\": HEIGHT_KD,\n",
        "    \"roll_kp\": ROLL_KP, \"roll_ki\": ROLL_KI, \"roll_kd\": ROLL_KD,\n",
        "    \"pitch_kp\": PITCH_KP, \"pitch_ki\": PITCH_KI, \"pitch_kd\": PITCH_KD,\n",
        "}\n",
        "\n",
        "print(\"PID Gains (used by simulation):\")\n",
        "print(f\"  Height: Kp={HEIGHT_KP}, Ki={HEIGHT_KI}, Kd={HEIGHT_KD}\")\n",
        "print(f\"  Roll:   Kp={ROLL_KP}, Ki={ROLL_KI}, Kd={ROLL_KD}\")\n",
        "print(f\"  Pitch:  Kp={PITCH_KP}, Ki={PITCH_KI}, Kd={PITCH_KD}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 95,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Loading maglev model from /Users/adipu/Documents/guadaloop_lev_control/RL Testing/maglev_model.pkl...\n",
            "Model loaded. Degree: 6\n",
            "Force R2: 1.0000\n",
            "Torque R2: 0.9999\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/Users/adipu/Documents/guadaloop_lev_control/.venv/lib/python3.13/site-packages/sklearn/base.py:463: InconsistentVersionWarning: Trying to unpickle estimator PolynomialFeatures from version 1.7.2 when using version 1.8.0. This might lead to breaking code or invalid results. Use at your own risk. For more info please refer to:\n",
            "https://scikit-learn.org/stable/model_persistence.html#security-maintainability-limitations\n",
            "  warnings.warn(\n",
            "/Users/adipu/Documents/guadaloop_lev_control/.venv/lib/python3.13/site-packages/sklearn/base.py:463: InconsistentVersionWarning: Trying to unpickle estimator LinearRegression from version 1.7.2 when using version 1.8.0. This might lead to breaking code or invalid results. Use at your own risk. For more info please refer to:\n",
            "https://scikit-learn.org/stable/model_persistence.html#security-maintainability-limitations\n",
            "  warnings.warn(\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Feedforward LUT ready. At target gap (11.86 mm): PWM = -0.0006\n"
          ]
        }
      ],
      "source": [
        "# ============================================================\n",
        "#           FEEDFORWARD USING MAGLEV PREDICTOR\n",
        "# ============================================================\n",
        "# Builds gap → equilibrium PWM LUT in pid_simulation (used by run_pid_simulation).\n",
        "build_feedforward_lut()\n",
        "print(f\"Feedforward LUT ready. At target gap ({TARGET_GAP*1000:.2f} mm): PWM = {feedforward_pwm(TARGET_GAP * 1000):.4f}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "---\n",
        "## Simulation Runner"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 96,
      "metadata": {},
      "outputs": [],
      "source": [
        "# run_pid_simulation is imported from pid_simulation (see imports).\n",
        "# It accepts gains=... (dict with height_kp, height_ki, ...), plus initial_gap_mm, max_steps, etc.\n",
        "# The gains dict is built in the \"PID Gains Configuration\" cell above."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "---\n",
        "## Plotting Functions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 97,
      "metadata": {},
      "outputs": [],
      "source": [
        "def plot_results(data: dict, title_suffix: str = \"\"):\n",
        "    \"\"\"\n",
        "    Create aesthetic plots of simulation results.\n",
        "    \n",
        "    Args:\n",
        "        data: Dictionary from run_pid_simulation()\n",
        "        title_suffix: Optional suffix for plot titles\n",
        "    \"\"\"\n",
        "    fig, axes = plt.subplots(2, 2, figsize=(14, 10))\n",
        "    fig.suptitle(f'PID Control Performance{title_suffix}', fontsize=16, fontweight='bold', y=1.02)\n",
        "    \n",
        "    target_gap_mm = TARGET_GAP * 1000\n",
        "    time = data['time']\n",
        "    \n",
        "    # Color palette\n",
        "    colors = {\n",
        "        'primary': '#2563eb',    # Blue\n",
        "        'secondary': '#7c3aed',  # Purple\n",
        "        'accent1': '#059669',    # Green\n",
        "        'accent2': '#dc2626',    # Red\n",
        "        'target': '#f59e0b',     # Orange\n",
        "        'grid': '#e5e7eb'\n",
        "    }\n",
        "    \n",
        "    # ========== Gap Height Plot ==========\n",
        "    ax1 = axes[0, 0]\n",
        "    ax1.plot(time, data['gap_avg'], color=colors['primary'], linewidth=2, label='Average Gap')\n",
        "    ax1.plot(time, data['gap_front'], color=colors['secondary'], linewidth=1, alpha=0.6, label='Front Yoke')\n",
        "    ax1.plot(time, data['gap_back'], color=colors['accent1'], linewidth=1, alpha=0.6, label='Back Yoke')\n",
        "    ax1.axhline(y=target_gap_mm, color=colors['target'], linestyle='--', linewidth=2, label=f'Target ({target_gap_mm:.1f}mm)')\n",
        "    \n",
        "    ax1.set_xlabel('Time (s)')\n",
        "    ax1.set_ylabel('Gap Height (mm)')\n",
        "    ax1.set_title('Gap Height Over Time', fontweight='bold')\n",
        "    ax1.legend(loc='best', framealpha=0.9)\n",
        "    ax1.set_ylim([0, 20])\n",
        "    ax1.grid(True, alpha=0.3)\n",
        "    \n",
        "    # Add settling info\n",
        "    final_error = abs(data['gap_avg'][-1] - target_gap_mm)\n",
        "    ax1.text(0.98, 0.02, f'Final error: {final_error:.2f}mm', \n",
        "             transform=ax1.transAxes, ha='right', va='bottom',\n",
        "             fontsize=10, bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))\n",
        "    \n",
        "    # ========== Roll Angle Plot ==========\n",
        "    ax2 = axes[0, 1]\n",
        "    ax2.plot(time, data['roll_deg'], color=colors['primary'], linewidth=2)\n",
        "    ax2.axhline(y=0, color=colors['target'], linestyle='--', linewidth=1.5, alpha=0.7)\n",
        "    ax2.fill_between(time, data['roll_deg'], 0, alpha=0.2, color=colors['primary'])\n",
        "    \n",
        "    ax2.set_xlabel('Time (s)')\n",
        "    ax2.set_ylabel('Roll Angle (degrees)')\n",
        "    ax2.set_title('Roll Angle Over Time', fontweight='bold')\n",
        "    ax2.grid(True, alpha=0.3)\n",
        "    \n",
        "    ax2.set_ylim([-4, 4])\n",
        "    \n",
        "    # ========== Pitch Angle Plot ==========\n",
        "    ax3 = axes[1, 0]\n",
        "    ax3.plot(time, data['pitch_deg'], color=colors['secondary'], linewidth=2)\n",
        "    ax3.axhline(y=0, color=colors['target'], linestyle='--', linewidth=1.5, alpha=0.7)\n",
        "    ax3.fill_between(time, data['pitch_deg'], 0, alpha=0.2, color=colors['secondary'])\n",
        "    \n",
        "    ax3.set_xlabel('Time (s)')\n",
        "    ax3.set_ylabel('Pitch Angle (degrees)')\n",
        "    ax3.set_title('Pitch Angle Over Time', fontweight='bold')\n",
        "    ax3.grid(True, alpha=0.3)\n",
        "    \n",
        "    ax3.set_ylim([-6, 6])\n",
        "    \n",
        "    # ========== Current Draw Plot ==========\n",
        "    ax4 = axes[1, 1]\n",
        "    ax4.plot(time, data['current_FL'], linewidth=1.5, label='Front Left', alpha=0.8)\n",
        "    ax4.plot(time, data['current_FR'], linewidth=1.5, label='Front Right', alpha=0.8)\n",
        "    ax4.plot(time, data['current_BL'], linewidth=1.5, label='Back Left', alpha=0.8)\n",
        "    ax4.plot(time, data['current_BR'], linewidth=1.5, label='Back Right', alpha=0.8)\n",
        "    ax4.plot(time, data['current_total'], color='black', linewidth=2, label='Total', linestyle='--')\n",
        "    \n",
        "    ax4.set_xlabel('Time (s)')\n",
        "    ax4.set_ylabel('Current (A)')\n",
        "    ax4.set_title('Coil Current Draw Over Time', fontweight='bold')\n",
        "    ax4.legend(loc='best', framealpha=0.9, ncol=2)\n",
        "    ax4.grid(True, alpha=0.3)\n",
        "    \n",
        "    # Add average power info\n",
        "    avg_total_current = np.mean(data['current_total'])\n",
        "    ax4.text(0.98, 0.98, f'Avg total: {avg_total_current:.2f}A', \n",
        "             transform=ax4.transAxes, ha='right', va='top',\n",
        "             fontsize=10, bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))\n",
        "    \n",
        "    plt.tight_layout()\n",
        "    plt.show()\n",
        "    \n",
        "    return fig\n",
        "\n",
        "\n",
        "def show_current_slice(data: dict, t_start: float, t_end: float):\n",
        "    \"\"\"\n",
        "    Plot coil current draw over a time slice (same as bottom-right graph in plot_results()).\n",
        "\n",
        "    Args:\n",
        "        data: Dictionary from run_pid_simulation()\n",
        "        t_start: Start time in seconds (inclusive)\n",
        "        t_end: End time in seconds (inclusive)\n",
        "    \"\"\"\n",
        "    time = np.asarray(data['time'])\n",
        "    mask = (time >= t_start) & (time <= t_end)\n",
        "    if not np.any(mask):\n",
        "        print(f'No data in range [{t_start}, {t_end}] s.')\n",
        "        return\n",
        "    t = time[mask]\n",
        "    fig, ax = plt.subplots(figsize=(10, 4))\n",
        "    ax.plot(t, np.asarray(data['current_FL'])[mask], linewidth=1.5, label='Front Left', alpha=0.8)\n",
        "    ax.plot(t, np.asarray(data['current_FR'])[mask], linewidth=1.5, label='Front Right', alpha=0.8)\n",
        "    ax.plot(t, np.asarray(data['current_BL'])[mask], linewidth=1.5, label='Back Left', alpha=0.8)\n",
        "    ax.plot(t, np.asarray(data['current_BR'])[mask], linewidth=1.5, label='Back Right', alpha=0.8)\n",
        "    ax.plot(t, np.asarray(data['current_total'])[mask], color='black', linewidth=2, label='Total', linestyle='--')\n",
        "    ax.set_xlabel('Time (s)')\n",
        "    ax.set_ylabel('Current (A)')\n",
        "    ax.set_title(f'Coil Current Draw ({t_start}–{t_end} s)', fontweight='bold')\n",
        "    ax.legend(loc='best', framealpha=0.9, ncol=2)\n",
        "    ax.grid(True, alpha=0.3)\n",
        "    avg_total = np.mean(np.asarray(data['current_total'])[mask])\n",
        "    ax.text(0.98, 0.98, f'Avg total: {avg_total:.2f}A', transform=ax.transAxes, ha='right', va='top',\n",
        "             fontsize=10, bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))\n",
        "    plt.tight_layout()\n",
        "    plt.show()\n",
        "    return fig"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 98,
      "metadata": {},
      "outputs": [],
      "source": [
        "def plot_control_signals(data: dict):\n",
        "    \"\"\"\n",
        "    Plot the PWM control signals sent to each coil.\n",
        "    \n",
        "    Args:\n",
        "        data: Dictionary from run_pid_simulation()\n",
        "    \"\"\"\n",
        "    fig, ax = plt.subplots(figsize=(12, 5))\n",
        "    \n",
        "    time = data['time']\n",
        "    \n",
        "    ax.plot(time, data['pwm_FL'], label='Front Left', linewidth=1.5, alpha=0.8)\n",
        "    ax.plot(time, data['pwm_FR'], label='Front Right', linewidth=1.5, alpha=0.8)\n",
        "    ax.plot(time, data['pwm_BL'], label='Back Left', linewidth=1.5, alpha=0.8)\n",
        "    ax.plot(time, data['pwm_BR'], label='Back Right', linewidth=1.5, alpha=0.8)\n",
        "    \n",
        "    if 'ff_pwm' in data and len(data.get('ff_pwm', [])) > 0:\n",
        "        ff = data['ff_pwm'] if isinstance(data['ff_pwm'], np.ndarray) else np.array(data['ff_pwm'])\n",
        "        ax.plot(time, ff, label='Feedforward', color='black', linewidth=2, linestyle='-.', alpha=0.7)\n",
        "    \n",
        "    ax.axhline(y=0, color='gray', linestyle='-', linewidth=0.5)\n",
        "    ax.axhline(y=1, color='red', linestyle='--', linewidth=1, alpha=0.5, label='Saturation')\n",
        "    ax.axhline(y=-1, color='red', linestyle='--', linewidth=1, alpha=0.5)\n",
        "    \n",
        "    ax.set_xlabel('Time (s)')\n",
        "    ax.set_ylabel('PWM Duty Cycle')\n",
        "    ax.set_title('Control Signals (PWM)', fontsize=14, fontweight='bold')\n",
        "    ax.legend(loc='best', framealpha=0.9)\n",
        "    ax.set_ylim([-1.2, 1.2])\n",
        "    ax.grid(True, alpha=0.3)\n",
        "    \n",
        "    plt.tight_layout()\n",
        "    plt.show()\n",
        "    \n",
        "    return fig"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "---\n",
        "## Run Simulation\n",
        "\n",
        "**Configure the simulation parameters below and run the cell.**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Loading maglev model from /Users/adipu/Documents/guadaloop_lev_control/RL Testing/maglev_model.pkl...\n",
            "Model loaded. Degree: 6\n",
            "Force R2: 1.0000\n",
            "Torque R2: 0.9999\n",
            "Starting simulation: initial_gap=9.0mm, target=11.86mm\n",
            "  Impulse disturbance: -2N at step 5000\n",
            "  Feedforward: enabled\n",
            "  Applied -2N impulse and 1.00 N·m torque at step 5000\n",
            "Simulation complete: 10000 steps, 41.66s\n",
            "  Final gap: 10.88mm (target: 11.86mm)\n",
            "  Final roll: 0.000°, pitch: 0.003°\n"
          ]
        }
      ],
      "source": [
        "# ============================================================\n",
        "#               SIMULATION PARAMETERS\n",
        "# ============================================================\n",
        "\n",
        "INITIAL_GAP_MM = 9.0     # Starting gap height (mm). Target is ~11.86mm\n",
        "MAX_STEPS = 10000          # Simulation steps (240 steps = 1 second)\n",
        "USE_GUI = False           # PyBullet GUI often hangs in notebooks; keep False. Use recording instead.\n",
        "USE_FEEDFORWARD = True    # Use MaglevPredictor feedforward for base PWM\n",
        "\n",
        "# Recording (headless: no GUI window; files saved to RECORD_DIR)\n",
        "RECORD_VIDEO = False      # Set True to save MP4 of each run (e.g. recordings/sim_YYYYMMDD_HHMMSS.mp4)\n",
        "RECORD_TELEMETRY = False  # Set True to save 4-panel telemetry PNG per run\n",
        "RECORD_DIR = \"recordings\" # Output folder for video and telemetry (created if missing)\n",
        "\n",
        "# Impulse disturbance (one-time force at a specific step)\n",
        "DISTURBANCE_STEP = 5000   # Step at which to apply impulse (e.g., 500). None = disabled\n",
        "DISTURBANCE_FORCE = -2   # Impulse force in Newtons (positive = upward push)\n",
        "\n",
        "# Stochastic disturbance (continuous random noise each step)\n",
        "DISTURBANCE_FORCE_STD = 0  # Std dev of random force noise (Newtons). 0 = disabled\n",
        "\n",
        "# ============================================================\n",
        "\n",
        "# Run simulation (gains from \"PID Gains Configuration\" cell)\n",
        "results = run_pid_simulation(\n",
        "    initial_gap_mm=INITIAL_GAP_MM,\n",
        "    max_steps=MAX_STEPS,\n",
        "    use_gui=USE_GUI,\n",
        "    disturbance_step=DISTURBANCE_STEP,\n",
        "    disturbance_force=DISTURBANCE_FORCE,\n",
        "    disturbance_force_std=DISTURBANCE_FORCE_STD,\n",
        "    use_feedforward=USE_FEEDFORWARD,\n",
        "    record_video=RECORD_VIDEO,\n",
        "    record_telemetry=RECORD_TELEMETRY,\n",
        "    record_dir=RECORD_DIR,\n",
        "    gains=gains,\n",
        "    verbose=True\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 100,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "{'time': array([0.00000000e+00, 4.16666667e-03, 8.33333333e-03, ...,\n",
            "       4.16541667e+01, 4.16583333e+01, 4.16625000e+01], shape=(10000,)), 'gap_front': array([ 9.        ,  9.02630146,  9.07173494, ..., 10.87328199,\n",
            "       10.88719472, 10.87328123], shape=(10000,)), 'gap_back': array([ 9.        ,  9.02630146,  9.07173494, ..., 10.88719405,\n",
            "       10.87328088, 10.88719392], shape=(10000,)), 'gap_avg': array([ 9.        ,  9.02630146,  9.07173494, ..., 10.88023802,\n",
            "       10.8802378 , 10.88023758], shape=(10000,)), 'roll_deg': array([-0.00000000e+00, -5.27915552e-06, -2.13414205e-05, ...,\n",
            "        1.30520451e-05,  1.30545699e-05,  1.30514835e-05], shape=(10000,)), 'pitch_deg': array([ 0.        ,  0.        ,  0.        , ...,  0.00316562,\n",
            "       -0.00316602,  0.00316576], shape=(10000,)), 'current_FL': array([10.2       ,  9.29838848,  8.99919415, ..., 10.2       ,\n",
            "       -3.9745388 , 10.2       ], shape=(10000,)), 'current_FR': array([10.2       ,  9.29830265,  8.9989996 , ..., 10.2       ,\n",
            "       -3.97454548, 10.2       ], shape=(10000,)), 'current_BL': array([10.2       ,  9.29838848,  8.99919415, ..., -3.97448826,\n",
            "       10.2       , -3.97456455], shape=(10000,)), 'current_BR': array([10.2       ,  9.29830265,  8.9989996 , ..., -3.97449398,\n",
            "       10.2       , -3.97457027], shape=(10000,)), 'current_total': array([40.8       , 37.19338226, 35.99638748, ..., 28.34898376,\n",
            "       28.34908295, 28.34913635], shape=(10000,)), 'pwm_FL': array([0.95018876, 0.8899194 , 0.83739257, ..., 0.4259686 , 0.22627316,\n",
            "       0.42596516], shape=(10000,), dtype=float32), 'pwm_FR': array([0.95018876, 0.8899151 , 0.83737934, ..., 0.4259683 , 0.22627287,\n",
            "       0.42596486], shape=(10000,), dtype=float32), 'pwm_BL': array([0.95018876, 0.8899194 , 0.83739257, ..., 0.22627567, 0.42596382,\n",
            "       0.22627184], shape=(10000,), dtype=float32), 'pwm_BR': array([0.95018876, 0.8899151 , 0.83737934, ..., 0.22627537, 0.42596352,\n",
            "       0.22627154], shape=(10000,), dtype=float32), 'ff_pwm': array([0.66143623, 0.65541619, 0.64500579, ..., 0.22719919, 0.22719919,\n",
            "       0.22719919], shape=(10000,))}\n"
          ]
        }
      ],
      "source": [
        "print(results)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "---\n",
        "## View Results"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 101,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "image/png": "[encoded data, 116688 characters]",
            "text/plain": [
              "<Figure size 1400x1000 with 4 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "image/png": "[encoded data, 39948 characters]",
            "text/plain": [
              "<Figure size 1000x400 with 1 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "image/png": "[encoded data, 39948 characters]",
            "text/plain": [
              "<Figure size 1000x400 with 1 Axes>"
            ]
          },
          "execution_count": 101,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# Plot main performance metrics\n",
        "plot_results(results)\n",
        "show_current_slice(results, 20, 25)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 102,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "image/png": "[encoded data, 43492 characters]",
            "text/plain": [
              "<Figure size 1200x500 with 1 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "image/png": "[encoded data, 43492 characters]",
            "text/plain": [
              "<Figure size 1200x500 with 1 Axes>"
            ]
          },
          "execution_count": 102,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# Plot control signals (optional)\n",
        "plot_control_signals(results)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "---\n",
        "## Compare Multiple Initial Conditions\n",
        "\n",
        "Run this cell to test the controller with different starting heights."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 103,
      "metadata": {},
      "outputs": [],
      "source": [
        "def compare_initial_conditions(initial_gaps_mm: list, max_steps: int = 2000):\n",
        "    \"\"\"\n",
        "    Compare PID performance across different initial conditions.\n",
        "    \n",
        "    Args:\n",
        "        initial_gaps_mm: List of starting gap heights to test\n",
        "        max_steps: Maximum steps per simulation\n",
        "    \"\"\"\n",
        "    fig, axes = plt.subplots(2, 2, figsize=(14, 10))\n",
        "    fig.suptitle('PID Performance Comparison: Different Initial Conditions', \n",
        "                 fontsize=16, fontweight='bold', y=1.02)\n",
        "    \n",
        "    target_gap_mm = TARGET_GAP * 1000\n",
        "    colors = plt.cm.viridis(np.linspace(0, 0.8, len(initial_gaps_mm)))\n",
        "    \n",
        "    all_results = []\n",
        "    \n",
        "    for i, gap in enumerate(initial_gaps_mm):\n",
        "        print(f\"Running simulation {i+1}/{len(initial_gaps_mm)}: initial_gap={gap}mm\")\n",
        "        data = run_pid_simulation(initial_gap_mm=gap, max_steps=max_steps, gains=gains, verbose=False)\n",
        "        all_results.append((gap, data))\n",
        "        \n",
        "        label = f'{gap}mm'\n",
        "        \n",
        "        # Gap height\n",
        "        axes[0, 0].plot(data['time'], data['gap_avg'], color=colors[i], \n",
        "                        linewidth=2, label=label)\n",
        "        \n",
        "        # Roll\n",
        "        axes[0, 1].plot(data['time'], data['roll_deg'], color=colors[i], \n",
        "                        linewidth=2, label=label)\n",
        "        \n",
        "        # Pitch\n",
        "        axes[1, 0].plot(data['time'], data['pitch_deg'], color=colors[i], \n",
        "                        linewidth=2, label=label)\n",
        "        \n",
        "        # Total current\n",
        "        axes[1, 1].plot(data['time'], data['current_total'], color=colors[i], \n",
        "                        linewidth=2, label=label)\n",
        "    \n",
        "    # Add target line to gap plot\n",
        "    axes[0, 0].axhline(y=target_gap_mm, color='red', linestyle='--', \n",
        "                       linewidth=2, label=f'Target ({target_gap_mm:.1f}mm)')\n",
        "    \n",
        "    # Configure axes\n",
        "    axes[0, 0].set_xlabel('Time (s)')\n",
        "    axes[0, 0].set_ylabel('Gap Height (mm)')\n",
        "    axes[0, 0].set_title('Average Gap Height', fontweight='bold')\n",
        "    axes[0, 0].legend(loc='best')\n",
        "    axes[0, 0].set_ylim([0, 20])\n",
        "    axes[0, 0].grid(True, alpha=0.3)\n",
        "    \n",
        "    axes[0, 1].axhline(y=0, color='gray', linestyle='--', linewidth=1)\n",
        "    axes[0, 1].set_xlabel('Time (s)')\n",
        "    axes[0, 1].set_ylabel('Roll Angle (degrees)')\n",
        "    axes[0, 1].set_title('Roll Angle', fontweight='bold')\n",
        "    axes[0, 1].legend(loc='best')\n",
        "    axes[0, 1].set_ylim([-4, 4])\n",
        "    axes[0, 1].grid(True, alpha=0.3)\n",
        "    \n",
        "    axes[1, 0].axhline(y=0, color='gray', linestyle='--', linewidth=1)\n",
        "    axes[1, 0].set_xlabel('Time (s)')\n",
        "    axes[1, 0].set_ylabel('Pitch Angle (degrees)')\n",
        "    axes[1, 0].set_title('Pitch Angle', fontweight='bold')\n",
        "    axes[1, 0].legend(loc='best')\n",
        "    axes[1, 0].set_ylim([-6, 6])\n",
        "    axes[1, 0].grid(True, alpha=0.3)\n",
        "    \n",
        "    axes[1, 1].set_xlabel('Time (s)')\n",
        "    axes[1, 1].set_ylabel('Total Current (A)')\n",
        "    axes[1, 1].set_title('Total Current Draw', fontweight='bold')\n",
        "    axes[1, 1].legend(loc='best')\n",
        "    axes[1, 1].grid(True, alpha=0.3)\n",
        "    \n",
        "    plt.tight_layout()\n",
        "    plt.show()\n",
        "    \n",
        "    return all_results"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 104,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "============================================================\n",
            "Data structure: results (from run_pid_simulation)\n",
            "============================================================\n",
            "Number of time steps: 10000\n",
            "\n",
            "Key                Shape          Dtype        Min        Max       \n",
            "------------------------------------------------------------\n",
            "current_BL         (10000,)       float64      -9.523     10.2      \n",
            "current_BR         (10000,)       float64      -9.365     10.2      \n",
            "current_FL         (10000,)       float64      -9.56      10.2      \n",
            "current_FR         (10000,)       float64      -9.274     10.2      \n",
            "current_total      (10000,)       float64      0.0007697  40.8      \n",
            "ff_pwm             (10000,)       float64      -0.003987  0.6614    \n",
            "gap_avg            (10000,)       float64      9          11.87     \n",
            "gap_back           (10000,)       float64      9          11.88     \n",
            "gap_front          (10000,)       float64      9          11.87     \n",
            "pitch_deg          (10000,)       float64      -0.004836  0.003869  \n",
            "pwm_BL             (10000,)       float32      -0.06742   0.9502    \n",
            "pwm_BR             (10000,)       float32      -0.05246   0.9502    \n",
            "pwm_FL             (10000,)       float32      -0.06965   0.9502    \n",
            "pwm_FR             (10000,)       float32      -0.05447   0.9502    \n",
            "roll_deg           (10000,)       float64      -0.1823    0.2749    \n",
            "time               (10000,)       float64      0          41.66     \n",
            "------------------------------------------------------------\n",
            "\n",
            "Sample (first and last time steps):\n",
            "\n",
            "  Step 0: t = 0.0000 s\n",
            "    gap_avg=9.000 mm  roll=-0.000°  pitch=0.000°\n",
            "    pwm: FL=0.950 FR=0.950 BL=0.950 BR=0.950\n",
            "    ff_pwm=0.661  current_total=40.800 A\n",
            "\n",
            "  Step 5000: t = 20.8333 s\n",
            "    gap_avg=11.860 mm  roll=0.002°  pitch=0.000°\n",
            "    pwm: FL=-0.001 FR=-0.001 BL=-0.001 BR=-0.001\n",
            "    ff_pwm=-0.001  current_total=0.019 A\n",
            "\n",
            "  Step 9999: t = 41.6625 s\n",
            "    gap_avg=10.880 mm  roll=0.000°  pitch=0.003°\n",
            "    pwm: FL=0.426 FR=0.426 BL=0.226 BR=0.226\n",
            "    ff_pwm=0.227  current_total=28.349 A\n"
          ]
        }
      ],
      "source": [
        "# Show simulation results data structure in a readable format\n",
        "def show_data_structure(data: dict, name: str = \"results\", sample_rows: int = 3):\n",
        "    \"\"\"Print keys, shapes, dtypes, and a small sample of the simulation data dict.\"\"\"\n",
        "    print(\"=\" * 60)\n",
        "    print(f\"Data structure: {name} (from run_pid_simulation)\")\n",
        "    print(\"=\" * 60)\n",
        "    n = len(data[\"time\"]) if \"time\" in data else 0\n",
        "    print(f\"Number of time steps: {n}\\n\")\n",
        "    print(f\"{'Key':<18} {'Shape':<14} {'Dtype':<12} {'Min':<10} {'Max':<10}\")\n",
        "    print(\"-\" * 60)\n",
        "    for key in sorted(data.keys()):\n",
        "        arr = data[key]\n",
        "        if isinstance(arr, np.ndarray):\n",
        "            sh = str(arr.shape)\n",
        "            dt = str(arr.dtype)\n",
        "            mn = f\"{np.min(arr):.4g}\" if arr.size else \"—\"\n",
        "            mx = f\"{np.max(arr):.4g}\" if arr.size else \"—\"\n",
        "        else:\n",
        "            sh, dt, mn, mx = \"—\", \"—\", \"—\", \"—\"\n",
        "        print(f\"{key:<18} {sh:<14} {dt:<12} {mn:<10} {mx:<10}\")\n",
        "    print(\"-\" * 60)\n",
        "    print(\"\\nSample (first and last time steps):\")\n",
        "    if n >= 1:\n",
        "        sample_idx = [0, n // 2, n - 1] if n >= 3 else list(range(n))\n",
        "        for i in sample_idx:\n",
        "            t = data[\"time\"][i]\n",
        "            print(f\"\\n  Step {i}: t = {t:.4f} s\")\n",
        "            print(f\"    gap_avg={data['gap_avg'][i]:.3f} mm  roll={data['roll_deg'][i]:.3f}°  pitch={data['pitch_deg'][i]:.3f}°\")\n",
        "            print(f\"    pwm: FL={data['pwm_FL'][i]:.3f} FR={data['pwm_FR'][i]:.3f} BL={data['pwm_BL'][i]:.3f} BR={data['pwm_BR'][i]:.3f}\")\n",
        "            if \"ff_pwm\" in data:\n",
        "                print(f\"    ff_pwm={data['ff_pwm'][i]:.3f}  current_total={data['current_total'][i]:.3f} A\")\n",
        "\n",
        "show_data_structure(results, \"results\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 105,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Compare different starting heights\n",
        "# Uncomment and run to test multiple initial conditions\n",
        "\n",
        "# comparison_results = compare_initial_conditions(\n",
        "#     initial_gaps_mm=[10.0, 14.0, 18.0, 22.0],\n",
        "#     max_steps=2000\n",
        "# )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "---\n",
        "## PID Tuning with Optuna\n",
        "\n",
        "Bayesian-style optimization (Optuna TPE sampler) tunes all **nine** gains (height, roll, pitch × Kp, Ki, Kd) **jointly** so loop coupling is respected.\n",
        "\n",
        "- **Run optimizer** (script or cell below): saves `pid_best_params.json` in this folder.\n",
        "- **Use new params**: re-run the **Load PID gains** cell at the top, then the **PID Gains** cell and the simulation."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 106,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Option A: Run Optuna optimization (slow). Saves pid_best_params.json next to optuna_pid_tune.py.\n",
        "# from optuna_pid_tune import run_optimization\n",
        "# study = run_optimization(n_trials=50, timeout=1800)\n",
        "# print(\"Best:\", study.best_params)\n",
        "\n",
        "# After running Optuna: re-run the \"Load PID gains\" cell at the top of the notebook\n",
        "# to pick up the new pid_best_params.json, then re-run the PID Gains cell and the simulation."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "---\n",
        "## Test Disturbance Rejection\n",
        "\n",
        "Apply a sudden force disturbance and observe recovery."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 107,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Test disturbance rejection\n",
        "# Uncomment and run to test disturbance response\n",
        "\n",
        "# Example 1: Impulse disturbance (one-time force)\n",
        "# disturbance_results = run_pid_simulation(\n",
        "#     initial_gap_mm=11.86,         # Start near target\n",
        "#     max_steps=3000,\n",
        "#     use_gui=False,\n",
        "#     disturbance_step=720,        # Apply at 3 seconds\n",
        "#     disturbance_force=-20.0,     # 20N downward push\n",
        "#     verbose=True\n",
        "# )\n",
        "# plot_results(disturbance_results, title_suffix=' (with 20N impulse at t=3s)')\n",
        "\n",
        "# Example 2: Continuous stochastic noise\n",
        "# noisy_results = run_pid_simulation(\n",
        "#     initial_gap_mm=11.86,\n",
        "#     max_steps=3000,\n",
        "#     use_gui=False,\n",
        "#     disturbance_force_std=2.0,   # 2N standard deviation continuous noise\n",
        "#     verbose=True\n",
        "# )\n",
        "# plot_results(noisy_results, title_suffix=' (with 2N std continuous noise)')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "---\n",
        "## Summary Statistics"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 108,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "==================================================\n",
            "SIMULATION SUMMARY\n",
            "==================================================\n",
            "Duration: 41.66 seconds (10000 steps)\n",
            "Target gap: 11.86 mm\n",
            "\n",
            "Gap Height:\n",
            "  Initial:  9.00 mm\n",
            "  Final:    10.88 mm\n",
            "  Error:    0.980 mm\n",
            "  Max:      11.87 mm\n",
            "  Min:      9.00 mm\n",
            "  Settling time (2%): 0.179 s\n",
            "\n",
            "Orientation (final):\n",
            "  Roll:  +0.000 degrees\n",
            "  Pitch: +0.003 degrees\n",
            "\n",
            "Current Draw:\n",
            "  Average total: 14.15 A\n",
            "  Peak total:    40.80 A\n",
            "==================================================\n"
          ]
        }
      ],
      "source": [
        "def print_summary(data: dict):\n",
        "    \"\"\"Print summary statistics for a simulation run.\"\"\"\n",
        "    target_gap_mm = TARGET_GAP * 1000\n",
        "    \n",
        "    # Calculate settling time (within 2% of target)\n",
        "    tolerance = 0.02 * target_gap_mm\n",
        "    settled_idx = None\n",
        "    for i in range(len(data['gap_avg'])):\n",
        "        if abs(data['gap_avg'][i] - target_gap_mm) < tolerance:\n",
        "            # Check if it stays settled\n",
        "            if all(abs(data['gap_avg'][j] - target_gap_mm) < tolerance \n",
        "                   for j in range(i, min(i+100, len(data['gap_avg'])))):\n",
        "                settled_idx = i\n",
        "                break\n",
        "    \n",
        "    print(\"=\" * 50)\n",
        "    print(\"SIMULATION SUMMARY\")\n",
        "    print(\"=\" * 50)\n",
        "    print(f\"Duration: {data['time'][-1]:.2f} seconds ({len(data['time'])} steps)\")\n",
        "    print(f\"Target gap: {target_gap_mm:.2f} mm\")\n",
        "    print()\n",
        "    print(\"Gap Height:\")\n",
        "    print(f\"  Initial:  {data['gap_avg'][0]:.2f} mm\")\n",
        "    print(f\"  Final:    {data['gap_avg'][-1]:.2f} mm\")\n",
        "    print(f\"  Error:    {abs(data['gap_avg'][-1] - target_gap_mm):.3f} mm\")\n",
        "    print(f\"  Max:      {max(data['gap_avg']):.2f} mm\")\n",
        "    print(f\"  Min:      {min(data['gap_avg']):.2f} mm\")\n",
        "    if settled_idx:\n",
        "        print(f\"  Settling time (2%): {data['time'][settled_idx]:.3f} s\")\n",
        "    else:\n",
        "        print(f\"  Settling time: Not settled within tolerance\")\n",
        "    print()\n",
        "    print(\"Orientation (final):\")\n",
        "    print(f\"  Roll:  {data['roll_deg'][-1]:+.3f} degrees\")\n",
        "    print(f\"  Pitch: {data['pitch_deg'][-1]:+.3f} degrees\")\n",
        "    print()\n",
        "    print(\"Current Draw:\")\n",
        "    print(f\"  Average total: {np.mean(data['current_total']):.2f} A\")\n",
        "    print(f\"  Peak total:    {max(data['current_total']):.2f} A\")\n",
        "    print(\"=\" * 50)\n",
        "\n",
        "# Print summary for last simulation\n",
        "print_summary(results)"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": ".venv",
      "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.13.9"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 4
}
