{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PID Control Testing for Maglev Pod\n",
    "\n",
    "This notebook provides a PID control interface for the `LevPodEnv` simulation environment.\n",
    "\n",
    "## Control Architecture\n",
    "- **Height PID**: Controls average gap height → base PWM 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 = height_output ± roll_adjustment ± pitch_adjustment`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from lev_pod_env import LevPodEnv, TARGET_GAP\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",
    "## PID Controller Class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PIDController:\n",
    "    \"\"\"Simple PID controller with anti-windup.\"\"\"\n",
    "    \n",
    "    def __init__(self, kp: float, ki: float, kd: float, \n",
    "                 output_min: float = -1.0, output_max: float = 1.0,\n",
    "                 integral_limit: float = np.inf):\n",
    "        self.kp = kp\n",
    "        self.ki = ki\n",
    "        self.kd = kd\n",
    "        self.output_min = output_min\n",
    "        self.output_max = output_max\n",
    "        self.integral_limit = integral_limit if integral_limit else abs(output_max) * 2\n",
    "        \n",
    "        self.integral = 0.0\n",
    "        self.prev_error = 0.0\n",
    "        self.first_update = True\n",
    "    \n",
    "    def reset(self):\n",
    "        \"\"\"Reset controller state.\"\"\"\n",
    "        self.integral = 0.0\n",
    "        self.prev_error = 0.0\n",
    "        self.first_update = True\n",
    "    \n",
    "    def update(self, error: float, dt: float) -> float:\n",
    "        \"\"\"Compute PID output.\n",
    "        \n",
    "        Args:\n",
    "            error: Current error (setpoint - measurement)\n",
    "            dt: Time step in seconds\n",
    "        \n",
    "        Returns:\n",
    "            Control output (clamped to output limits)\n",
    "        \"\"\"\n",
    "        # Proportional term\n",
    "        p_term = self.kp * error\n",
    "        \n",
    "        # Integral term with anti-windup\n",
    "        self.integral += error * dt\n",
    "        self.integral = np.clip(self.integral, -self.integral_limit, self.integral_limit)\n",
    "        i_term = self.ki * self.integral\n",
    "        \n",
    "        # Derivative term (skip on first update to avoid spike)\n",
    "        if self.first_update:\n",
    "            d_term = 0.0\n",
    "            self.first_update = False\n",
    "        else:\n",
    "            d_term = self.kd * (error - self.prev_error) / dt\n",
    "        \n",
    "        self.prev_error = error\n",
    "        \n",
    "        # Compute and clamp output\n",
    "        output = p_term + i_term + d_term\n",
    "        return np.clip(output, self.output_min, self.output_max)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## PID Gains Configuration\n",
    "\n",
    "**Modify these values to tune the controller.**\n",
    "\n",
    "- **Height PID**: Controls vertical position (gap height)\n",
    "- **Roll PID**: Corrects side-to-side tilt\n",
    "- **Pitch PID**: Corrects front-to-back tilt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ============================================================\n",
    "#                    HEIGHT PID GAINS\n",
    "# ============================================================\n",
    "# Controls average gap height. Positive error = pod too low.\n",
    "# Output: base PWM applied to all coils (positive = more lift)\n",
    "\n",
    "HEIGHT_KP = 50.0    # Proportional gain\n",
    "HEIGHT_KI = 5.0     # Integral gain  \n",
    "HEIGHT_KD = 10.0    # Derivative gain\n",
    "\n",
    "# ============================================================\n",
    "#                     ROLL PID GAINS\n",
    "# ============================================================\n",
    "# Corrects roll angle (rotation about X-axis).\n",
    "# Positive roll = right side lower → increase right coils\n",
    "# Output: differential adjustment to left/right coils\n",
    "\n",
    "ROLL_KP = 2.0       # Proportional gain\n",
    "ROLL_KI = 0.5       # Integral gain\n",
    "ROLL_KD = 0.5       # Derivative gain\n",
    "\n",
    "# ============================================================\n",
    "#                    PITCH PID GAINS  \n",
    "# ============================================================\n",
    "# Corrects pitch angle (rotation about Y-axis).\n",
    "# Positive pitch = back lower → increase back coils\n",
    "# Output: differential adjustment to front/back coils\n",
    "\n",
    "PITCH_KP = 2.0      # Proportional gain\n",
    "PITCH_KI = 0.5      # Integral gain\n",
    "PITCH_KD = 0.5      # Derivative gain\n",
    "\n",
    "print(\"PID Gains Configured:\")\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": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Simulation Runner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_pid_simulation(\n",
    "    initial_gap_mm: float = 14.0,\n",
    "    max_steps: int = 2000,\n",
    "    use_gui: bool = False,\n",
    "    disturbance_step: int = None,\n",
    "    disturbance_force: float = 0.0,\n",
    "    disturbance_force_std: float = 0.0,\n",
    "    verbose: bool = True\n",
    ") -> dict:\n",
    "    \"\"\"\n",
    "    Run a PID-controlled simulation of the maglev pod.\n",
    "\n",
    "    Args:\n",
    "        initial_gap_mm: Starting gap height in millimeters\n",
    "        max_steps: Maximum simulation steps (each step is 1/240 second)\n",
    "        use_gui: Whether to show PyBullet GUI visualization\n",
    "        disturbance_step: Step at which to apply impulse disturbance (None = no impulse)\n",
    "        disturbance_force: Impulse force in Newtons at disturbance_step (positive = upward)\n",
    "        disturbance_force_std: Standard deviation of continuous stochastic noise (Newtons)\n",
    "        verbose: Print progress updates\n",
    "\n",
    "    Returns:\n",
    "        Dictionary containing time series data for plotting\n",
    "    \"\"\"\n",
    "    # Create environment with stochastic disturbance if specified\n",
    "    env = LevPodEnv(\n",
    "        use_gui=use_gui,\n",
    "        initial_gap_mm=initial_gap_mm,\n",
    "        max_steps=max_steps,\n",
    "        disturbance_force_std=disturbance_force_std\n",
    "    )\n",
    "    dt = env.dt  # 1/240 seconds\n",
    "\n",
    "    # Initialize PID controllers\n",
    "    height_pid = PIDController(HEIGHT_KP, HEIGHT_KI, HEIGHT_KD, output_min=-1.0, output_max=1.0)\n",
    "    roll_pid = PIDController(ROLL_KP, ROLL_KI, ROLL_KD, output_min=-0.5, output_max=0.5)\n",
    "    pitch_pid = PIDController(PITCH_KP, PITCH_KI, PITCH_KD, output_min=-0.5, output_max=0.5)\n",
    "\n",
    "    # Data storage\n",
    "    data = {\n",
    "        'time': [],\n",
    "        'gap_front': [],\n",
    "        'gap_back': [],\n",
    "        'gap_avg': [],\n",
    "        'roll_deg': [],\n",
    "        'pitch_deg': [],\n",
    "        'current_FL': [],\n",
    "        'current_FR': [],\n",
    "        'current_BL': [],\n",
    "        'current_BR': [],\n",
    "        'current_total': [],\n",
    "        'pwm_FL': [],\n",
    "        'pwm_FR': [],\n",
    "        'pwm_BL': [],\n",
    "        'pwm_BR': [],\n",
    "    }\n",
    "\n",
    "    # Reset environment\n",
    "    obs, _ = env.reset()\n",
    "\n",
    "    target_gap = TARGET_GAP  # meters\n",
    "\n",
    "    if verbose:\n",
    "        print(f\"Starting simulation: initial_gap={initial_gap_mm}mm, target={target_gap*1000:.2f}mm\")\n",
    "        if disturbance_step:\n",
    "            print(f\"  Impulse disturbance: {disturbance_force}N at step {disturbance_step}\")\n",
    "        if disturbance_force_std > 0:\n",
    "            print(f\"  Stochastic noise: std={disturbance_force_std}N\")\n",
    "\n",
    "    for step in range(max_steps):\n",
    "        t = step * dt\n",
    "\n",
    "        # Extract gap heights from observation (first 4 values are normalized gaps)\n",
    "        # These are sensor readings (may include noise in future)\n",
    "        gaps_normalized = obs[:4]  # [center_right, center_left, front, back]\n",
    "        gaps = gaps_normalized * env.gap_scale + TARGET_GAP\n",
    "\n",
    "        gap_front = gaps[2]  # Front sensor\n",
    "        gap_back = gaps[3]   # Back sensor\n",
    "        gap_left = gaps[1]   # Center_Left sensor\n",
    "        gap_right = gaps[0]  # Center_Right sensor\n",
    "\n",
    "        gap_avg = (gap_front + gap_back + gap_left + gap_right) / 4\n",
    "\n",
    "        # Calculate roll and pitch from sensor gaps (for PID control)\n",
    "        y_distance = 0.1016  # Left to right distance (m)\n",
    "        x_distance = 0.2518  # Front to back distance (m)\n",
    "\n",
    "        roll_angle = np.arcsin(np.clip((gap_left - gap_right) / y_distance, -1, 1))\n",
    "        pitch_angle = np.arcsin(np.clip((gap_back - gap_front) / x_distance, -1, 1))\n",
    "\n",
    "        # Compute PID outputs\n",
    "        height_error = target_gap - gap_avg  # Positive = too low, need more lift\n",
    "        roll_error = -roll_angle  # Target is 0, negative feedback\n",
    "        pitch_error = -pitch_angle  # Target is 0, negative feedback\n",
    "\n",
    "        height_output = height_pid.update(height_error, dt)\n",
    "        roll_output = roll_pid.update(roll_error, dt)\n",
    "        pitch_output = pitch_pid.update(pitch_error, dt)\n",
    "\n",
    "        # Combine outputs into 4 coil PWM commands\n",
    "        # Coil layout: front_left (+x,+y), front_right (+x,-y), back_left (-x,+y), back_right (-x,-y)\n",
    "        # Roll correction: positive roll_output → increase right side (FR, BR)\n",
    "        # Pitch correction: positive pitch_output → increase back side (BL, BR)\n",
    "\n",
    "        pwm_FL = np.clip(height_output - roll_output - pitch_output, -1, 1)\n",
    "        pwm_FR = np.clip(height_output + roll_output - pitch_output, -1, 1)\n",
    "        pwm_BL = np.clip(height_output - roll_output + pitch_output, -1, 1)\n",
    "        pwm_BR = np.clip(height_output + roll_output + pitch_output, -1, 1)\n",
    "\n",
    "        action = np.array([pwm_FL, pwm_FR, pwm_BL, pwm_BR], dtype=np.float32)\n",
    "\n",
    "        # Apply impulse disturbance if specified (uses environment method)\n",
    "        if disturbance_step and step == disturbance_step:\n",
    "            env.apply_impulse(disturbance_force)\n",
    "            if verbose:\n",
    "                print(f\"  Applied {disturbance_force}N impulse at step {step}\")\n",
    "\n",
    "        # Step environment (stochastic disturbance applied internally if enabled)\n",
    "        obs, _, terminated, truncated, info = env.step(action)\n",
    "\n",
    "        # Store ground truth data from environment info (not sensor observations)\n",
    "        data['time'].append(t)\n",
    "        data['gap_front'].append(info['gap_front_yoke'] * 1000)  # Convert to mm\n",
    "        data['gap_back'].append(info['gap_back_yoke'] * 1000)\n",
    "        data['gap_avg'].append(info['gap_avg'] * 1000)  # Use ground truth from info\n",
    "        data['roll_deg'].append(np.degrees(info['roll']))\n",
    "        data['pitch_deg'].append(np.degrees(info['pitch']))  # Use ground truth from info\n",
    "        data['current_FL'].append(info['curr_front_L'])\n",
    "        data['current_FR'].append(info['curr_front_R'])\n",
    "        data['current_BL'].append(info['curr_back_L'])\n",
    "        data['current_BR'].append(info['curr_back_R'])\n",
    "        data['current_total'].append(\n",
    "            abs(info['curr_front_L']) + abs(info['curr_front_R']) +\n",
    "            abs(info['curr_back_L']) + abs(info['curr_back_R'])\n",
    "        )\n",
    "        data['pwm_FL'].append(pwm_FL)\n",
    "        data['pwm_FR'].append(pwm_FR)\n",
    "        data['pwm_BL'].append(pwm_BL)\n",
    "        data['pwm_BR'].append(pwm_BR)\n",
    "\n",
    "        if terminated or truncated:\n",
    "            if verbose:\n",
    "                print(f\"  Simulation ended at step {step} (terminated={terminated})\")\n",
    "            break\n",
    "\n",
    "    env.close()\n",
    "\n",
    "    # Convert to numpy arrays\n",
    "    for key in data:\n",
    "        data[key] = np.array(data[key])\n",
    "\n",
    "    if verbose:\n",
    "        print(f\"Simulation complete: {len(data['time'])} steps, {data['time'][-1]:.2f}s\")\n",
    "        print(f\"  Final gap: {data['gap_avg'][-1]:.2f}mm (target: {target_gap*1000:.2f}mm)\")\n",
    "        print(f\"  Final roll: {data['roll_deg'][-1]:.3f}°, pitch: {data['pitch_deg'][-1]:.3f}°\")\n",
    "\n",
    "    return data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Plotting Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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, max(35, max(data['gap_avg']) * 1.1)])\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",
    "    max_roll = max(abs(data['roll_deg'].min()), abs(data['roll_deg'].max()))\n",
    "    ax2.set_ylim([-max(1, max_roll * 1.2), max(1, max_roll * 1.2)])\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",
    "    max_pitch = max(abs(data['pitch_deg'].min()), abs(data['pitch_deg'].max()))\n",
    "    ax3.set_ylim([-max(1, max_pitch * 1.2), max(1, max_pitch * 1.2)])\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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",
    "    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": [],
   "source": [
    "# ============================================================\n",
    "#               SIMULATION PARAMETERS\n",
    "# ============================================================\n",
    "\n",
    "INITIAL_GAP_MM = 14.0     # Starting gap height (mm). Target is ~16.49mm\n",
    "MAX_STEPS = 2000          # Simulation steps (240 steps = 1 second)\n",
    "USE_GUI = False           # Set to True to see PyBullet visualization\n",
    "\n",
    "# Impulse disturbance (one-time force at a specific step)\n",
    "DISTURBANCE_STEP = None   # Step at which to apply impulse (e.g., 500). None = disabled\n",
    "DISTURBANCE_FORCE = 0.0   # Impulse force in Newtons (positive = upward push)\n",
    "\n",
    "# Stochastic disturbance (continuous random noise each step)\n",
    "DISTURBANCE_FORCE_STD = 0.0  # Std dev of random force noise (Newtons). 0 = disabled\n",
    "\n",
    "# ============================================================\n",
    "\n",
    "# Run simulation\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",
    "    verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## View Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot main performance metrics\n",
    "plot_results(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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": null,
   "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, 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].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].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].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": null,
   "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",
    "## Test Disturbance Rejection\n",
    "\n",
    "Apply a sudden force disturbance and observe recovery."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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=16.5,         # 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=16.5,\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": null,
   "metadata": {},
   "outputs": [],
   "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": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}