"""
Tesla Coil Spark Course - Image Generation Script

Generates all course images programmatically using matplotlib, schemdraw, and numpy.
Run this script from the spark-lessons directory.

Usage:
    python generate_images.py                    # Generate all images
    python generate_images.py --part 1           # Generate Part 1 images only
    python generate_images.py --category graphs  # Generate only graphs
"""

import matplotlib.pyplot as plt
import matplotlib.patches as patches
from matplotlib.patches import FancyBboxPatch, FancyArrowPatch, Circle, Rectangle, Wedge
import numpy as np
import schemdraw
import schemdraw.elements as elm
from pathlib import Path
import argparse

# ============================================================================
# CONFIGURATION
# ============================================================================

# Common styling
STYLE = {
    'dpi': 150,
    'figure_facecolor': 'white',
    'text_color': '#2C3E50',
    'primary_color': '#3498DB',      # Blue
    'secondary_color': '#E74C3C',    # Red
    'accent_color': '#2ECC71',       # Green
    'warning_color': '#F39C12',      # Orange
    'grid_alpha': 0.3,
    'font_family': 'sans-serif',
    'title_size': 14,
    'label_size': 12,
    'tick_size': 10,
    'legend_size': 10,
}

# Directories
BASE_DIR = Path(__file__).parent
ASSETS_DIRS = {
    'fundamentals': BASE_DIR / 'lessons' / '01-fundamentals' / 'assets',
    'optimization': BASE_DIR / 'lessons' / '02-optimization' / 'assets',
    'spark-physics': BASE_DIR / 'lessons' / '03-spark-physics' / 'assets',
    'advanced-modeling': BASE_DIR / 'lessons' / '04-advanced-modeling' / 'assets',
    'shared': BASE_DIR / 'assets' / 'shared',
}

# Create directories if they don't exist
for dir_path in ASSETS_DIRS.values():
    dir_path.mkdir(parents=True, exist_ok=True)

# ============================================================================
# UTILITY FUNCTIONS
# ============================================================================

def set_style(ax):
    """Apply common styling to a matplotlib axis"""
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.tick_params(labelsize=STYLE['tick_size'])
    ax.grid(True, alpha=STYLE['grid_alpha'], linestyle='--')


def save_figure(fig, filename, directory='fundamentals'):
    """Save figure to appropriate directory"""
    filepath = ASSETS_DIRS[directory] / filename
    fig.savefig(filepath, dpi=STYLE['dpi'], bbox_inches='tight', facecolor='white')
    plt.close(fig)
    print(f"[OK] Generated: {filepath}")


# ============================================================================
# PART 1: FUNDAMENTALS IMAGES
# ============================================================================

def generate_complex_plane_admittance():
    """Image 3: Complex plane showing Y and Z phasors"""
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))

    # Left: Admittance plane
    Y_real = 10  # mS
    Y_imag = 15  # mS

    ax1.arrow(0, 0, Y_real, 0, head_width=1, head_length=0.8, fc=STYLE['primary_color'], ec=STYLE['primary_color'], linewidth=2)
    ax1.arrow(0, 0, 0, Y_imag, head_width=0.8, head_length=1, fc=STYLE['secondary_color'], ec=STYLE['secondary_color'], linewidth=2)
    ax1.arrow(0, 0, Y_real, Y_imag, head_width=1, head_length=1, fc='black', ec='black', linewidth=2.5)

    ax1.text(Y_real/2, -2, 'Re{Y} = G\n(Conductance)', ha='center', fontsize=STYLE['label_size'], color=STYLE['primary_color'])
    ax1.text(-2, Y_imag/2, 'Im{Y} = B\n(Susceptance)', ha='center', fontsize=STYLE['label_size'], color=STYLE['secondary_color'], rotation=90)
    ax1.text(Y_real/2 + 1, Y_imag/2 + 1, f'Y = {Y_real}+j{Y_imag} mS', fontsize=STYLE['label_size'], fontweight='bold')

    # Draw angle
    angle = np.arctan2(Y_imag, Y_real)
    arc = Wedge((0, 0), 3, 0, np.degrees(angle), facecolor='yellow', alpha=0.3, edgecolor='black')
    ax1.add_patch(arc)
    ax1.text(4, 2, f'θ_Y = {np.degrees(angle):.1f}°', fontsize=STYLE['label_size'])

    ax1.set_xlim(-5, 20)
    ax1.set_ylim(-5, 20)
    ax1.set_xlabel('Real Axis (mS)', fontsize=STYLE['label_size'])
    ax1.set_ylabel('Imaginary Axis (mS)', fontsize=STYLE['label_size'])
    ax1.set_title('Admittance (Y) Plane', fontsize=STYLE['title_size'], fontweight='bold')
    ax1.axhline(y=0, color='k', linewidth=0.5)
    ax1.axvline(x=0, color='k', linewidth=0.5)
    ax1.grid(True, alpha=STYLE['grid_alpha'])
    ax1.set_aspect('equal')

    # Right: Impedance plane
    Z_real = 30  # Ω
    Z_imag = -45  # Ω (capacitive)

    ax2.arrow(0, 0, Z_real, 0, head_width=3, head_length=2, fc=STYLE['primary_color'], ec=STYLE['primary_color'], linewidth=2)
    ax2.arrow(0, 0, 0, Z_imag, head_width=2, head_length=3, fc=STYLE['secondary_color'], ec=STYLE['secondary_color'], linewidth=2)
    ax2.arrow(0, 0, Z_real, Z_imag, head_width=3, head_length=3, fc='black', ec='black', linewidth=2.5)

    ax2.text(Z_real/2, 5, 'Re{Z} = R\n(Resistance)', ha='center', fontsize=STYLE['label_size'], color=STYLE['primary_color'])
    ax2.text(-5, Z_imag/2, 'Im{Z} = X\n(Reactance)', ha='center', fontsize=STYLE['label_size'], color=STYLE['secondary_color'], rotation=90)
    ax2.text(Z_real/2 + 3, Z_imag/2 - 5, f'Z = {Z_real}{Z_imag:+}j Ω', fontsize=STYLE['label_size'], fontweight='bold')

    # Draw angle
    angle_Z = np.arctan2(Z_imag, Z_real)
    arc2 = Wedge((0, 0), 8, np.degrees(angle_Z), 0, facecolor='lightblue', alpha=0.3, edgecolor='black')
    ax2.add_patch(arc2)
    ax2.text(10, -8, f'φ_Z = {np.degrees(angle_Z):.1f}°', fontsize=STYLE['label_size'])

    ax2.text(Z_real/2, Z_imag - 10, 'Capacitive\n(φ_Z < 0)', ha='center', fontsize=10, style='italic', color=STYLE['secondary_color'])

    ax2.set_xlim(-15, 50)
    ax2.set_ylim(-60, 15)
    ax2.set_xlabel('Real Axis (Ω)', fontsize=STYLE['label_size'])
    ax2.set_ylabel('Imaginary Axis (Ω)', fontsize=STYLE['label_size'])
    ax2.set_title('Impedance (Z) Plane', fontsize=STYLE['title_size'], fontweight='bold')
    ax2.axhline(y=0, color='k', linewidth=0.5)
    ax2.axvline(x=0, color='k', linewidth=0.5)
    ax2.grid(True, alpha=STYLE['grid_alpha'])
    ax2.set_aspect('equal')

    fig.suptitle('Complex Plane Representation: Admittance vs Impedance\nNote: φ_Z = -θ_Y',
                 fontsize=STYLE['title_size']+2, fontweight='bold', y=1.02)

    save_figure(fig, 'complex-plane-admittance.png', 'fundamentals')


def generate_phase_angle_visualization():
    """Image 4: Impedance phasors showing different phase angles"""
    fig, ax = plt.subplots(figsize=(10, 8))

    # Define impedances with different phase angles
    impedances = [
        (100, 0, '0°', 'Pure Resistive', STYLE['primary_color']),
        (100, -58, '-30°', 'Slightly Capacitive', STYLE['accent_color']),
        (71, -71, '-45°', 'Balanced (often unachievable)', STYLE['warning_color']),
        (50, -87, '-60°', 'More Capacitive', 'purple'),
        (26, -97, '-75°', 'Highly Capacitive (typical spark)', STYLE['secondary_color']),
    ]

    for Z_real, Z_imag, angle_label, description, color in impedances:
        # Draw phasor
        ax.arrow(0, 0, Z_real, Z_imag, head_width=5, head_length=4,
                fc=color, ec=color, linewidth=2, alpha=0.7, length_includes_head=True)

        # Label
        offset_x = 5 if Z_real > 50 else -15
        offset_y = -5
        ax.text(Z_real + offset_x, Z_imag + offset_y,
               f'{angle_label}\n{description}',
               fontsize=9, ha='left' if Z_real > 50 else 'right', color=color, fontweight='bold')

        # Power factor
        pf = Z_real / np.sqrt(Z_real**2 + Z_imag**2)
        ax.text(Z_real + offset_x, Z_imag + offset_y - 8,
               f'PF = {pf:.3f}',
               fontsize=8, ha='left' if Z_real > 50 else 'right', color=color, style='italic')

    # Highlight typical spark range
    theta1 = np.radians(-75)
    theta2 = np.radians(-55)
    r = 110
    angles = np.linspace(theta1, theta2, 50)
    x = r * np.cos(angles)
    y = r * np.sin(angles)
    ax.fill(np.concatenate([[0], x, [0]]), np.concatenate([[0], y, [0]]),
            alpha=0.2, color='lightcoral', label='Typical Tesla Coil Spark Range')

    # Add note about -45°
    ax.annotate('Often mathematically\nimpossible for sparks!',
                xy=(71, -71), xytext=(80, -40),
                arrowprops=dict(arrowstyle='->', color='red', lw=2),
                fontsize=10, color='red', fontweight='bold',
                bbox=dict(boxstyle='round,pad=0.5', facecolor='yellow', alpha=0.7))

    ax.set_xlim(-20, 120)
    ax.set_ylim(-110, 20)
    ax.set_xlabel('Re{Z} = Resistance (kΩ)', fontsize=STYLE['label_size'])
    ax.set_ylabel('Im{Z} = Reactance (kΩ)', fontsize=STYLE['label_size'])
    ax.set_title('Impedance Phase Angles and Their Physical Meanings',
                fontsize=STYLE['title_size'], fontweight='bold')
    ax.axhline(y=0, color='k', linewidth=0.8)
    ax.axvline(x=0, color='k', linewidth=0.8)
    ax.grid(True, alpha=STYLE['grid_alpha'])
    ax.legend(loc='upper right', fontsize=STYLE['legend_size'])
    ax.set_aspect('equal')

    save_figure(fig, 'phase-angle-visualization.png', 'fundamentals')


def generate_phase_constraint_graph():
    """Image 5: Graph of minimum achievable phase angle vs capacitance ratio"""
    fig, ax = plt.subplots(figsize=(10, 7))

    # Calculate φ_Z,min vs r
    r = np.linspace(0, 3, 300)
    phi_Z_min = -np.degrees(np.arctan(2 * np.sqrt(r * (1 + r))))

    # Plot curve
    ax.plot(r, phi_Z_min, linewidth=3, color=STYLE['primary_color'], label='φ_Z,min = -atan(2√[r(1+r)])')

    # Mark critical point r = 0.207
    r_crit = 0.207
    phi_crit = -np.degrees(np.arctan(2 * np.sqrt(r_crit * (1 + r_crit))))
    ax.plot(r_crit, phi_crit, 'ro', markersize=12, label=f'Critical: r = {r_crit:.3f}, φ_Z,min = -45°')

    # Shade impossible region
    ax.fill_between(r, phi_Z_min, -45, alpha=0.3, color='red', label='Impossible Region')

    # Add -45° line
    ax.axhline(y=-45, color='green', linestyle='--', linewidth=2, label='Traditional "Matched" Target (-45°)')

    # Shade typical Tesla coil region
    ax.axvspan(0.5, 2.0, alpha=0.2, color='yellow', label='Typical Tesla Coil Range')

    # Annotations for geometric examples
    examples = [
        (0.5, 'Short topload,\nlong spark'),
        (1.0, 'Balanced\ngeometry'),
        (2.0, 'Large topload,\nshort spark'),
    ]

    for r_ex, text in examples:
        phi_ex = -np.degrees(np.arctan(2 * np.sqrt(r_ex * (1 + r_ex))))
        ax.plot(r_ex, phi_ex, 'ks', markersize=8)
        ax.annotate(text, xy=(r_ex, phi_ex), xytext=(r_ex + 0.3, phi_ex - 5),
                   fontsize=9, ha='left',
                   arrowprops=dict(arrowstyle='->', color='black', lw=1))

    ax.set_xlim(0, 3)
    ax.set_ylim(-90, 0)
    ax.set_xlabel('r = C_mut / C_sh', fontsize=STYLE['label_size'])
    ax.set_ylabel('Minimum Impedance Phase φ_Z,min (degrees)', fontsize=STYLE['label_size'])
    ax.set_title('Topological Phase Constraint: Minimum Achievable Phase Angle',
                fontsize=STYLE['title_size'], fontweight='bold')
    ax.grid(True, alpha=STYLE['grid_alpha'])
    ax.legend(loc='lower left', fontsize=STYLE['legend_size'] - 1)

    # Add text box with key insight
    textstr = 'Key Insight:\nWhen r ≥ 0.207, achieving -45° is\nmathematically impossible regardless\nof resistance value!'
    props = dict(boxstyle='round', facecolor='wheat', alpha=0.8)
    ax.text(2.3, -20, textstr, fontsize=10, verticalalignment='top', bbox=props)

    save_figure(fig, 'phase-constraint-graph.png', 'fundamentals')


def generate_admittance_vector_addition():
    """Image 7: Vector diagram showing parallel admittance addition"""
    fig, ax = plt.subplots(figsize=(8, 8))

    # Branch 1: Y1 = G + jB1
    G = 8
    B1 = 12
    ax.arrow(0, 0, G, B1, head_width=0.8, head_length=0.6,
            fc=STYLE['primary_color'], ec=STYLE['primary_color'], linewidth=2.5, label='Y₁ = G + jB₁')
    ax.text(G/2 - 2, B1/2 + 1, 'Y₁ = G + jB₁\n(R || C_mut)', fontsize=11, color=STYLE['primary_color'], fontweight='bold')

    # Branch 2: Y2 = jB2
    B2 = 6
    ax.arrow(0, 0, 0, B2, head_width=0.6, head_length=0.5,
            fc=STYLE['secondary_color'], ec=STYLE['secondary_color'], linewidth=2.5, label='Y₂ = jB₂')
    ax.text(1, B2/2, 'Y₂ = jB₂\n(C_sh)', fontsize=11, color=STYLE['secondary_color'], fontweight='bold')

    # Total: Y_total = Y1 + Y2
    Y_total_real = G
    Y_total_imag = B1 + B2
    ax.arrow(0, 0, Y_total_real, Y_total_imag, head_width=1, head_length=0.8,
            fc='black', ec='black', linewidth=3, label='Y_total = Y₁ + Y₂', zorder=10)
    ax.text(Y_total_real/2 + 2, Y_total_imag/2, 'Y_total\n= G + j(B₁+B₂)',
           fontsize=12, fontweight='bold', bbox=dict(boxstyle='round', facecolor='yellow', alpha=0.7))

    # Draw parallelogram construction lines
    ax.plot([G, G], [B1, Y_total_imag], 'k--', alpha=0.4, linewidth=1)
    ax.plot([0, G], [B2, Y_total_imag], 'k--', alpha=0.4, linewidth=1)

    # Draw components
    ax.plot([0, G], [0, 0], 'b:', linewidth=2, alpha=0.5)
    ax.text(G/2, -0.8, 'G (conductance)', fontsize=10, ha='center', color='blue')

    ax.plot([0, 0], [0, B1], 'r:', linewidth=2, alpha=0.5)
    ax.text(-1.2, B1/2, 'B₁', fontsize=10, ha='center', color='red')

    ax.plot([0, 0], [B1, Y_total_imag], 'r:', linewidth=2, alpha=0.5)
    ax.text(-1.2, (B1 + Y_total_imag)/2, 'B₂', fontsize=10, ha='center', color='red')

    ax.set_xlim(-3, 15)
    ax.set_ylim(-2, 22)
    ax.set_xlabel('Real Part: G = 1/R (mS)', fontsize=STYLE['label_size'])
    ax.set_ylabel('Imaginary Part: B = ωC (mS)', fontsize=STYLE['label_size'])
    ax.set_title('Parallel Admittance Addition (Vector/Phasor Method)',
                fontsize=STYLE['title_size'], fontweight='bold')
    ax.axhline(y=0, color='k', linewidth=0.8)
    ax.axvline(x=0, color='k', linewidth=0.8)
    ax.grid(True, alpha=STYLE['grid_alpha'])
    ax.legend(loc='upper left', fontsize=STYLE['legend_size'])
    ax.set_aspect('equal')

    # Add formula
    formula_text = 'For parallel branches:\nY_total = Y₁ + Y₂ + Y₃ + ...\n(Admittances add directly!)'
    ax.text(10, 3, formula_text, fontsize=10,
           bbox=dict(boxstyle='round', facecolor='lightblue', alpha=0.7))

    save_figure(fig, 'admittance-vector-addition.png', 'fundamentals')


# ============================================================================
# PART 2: OPTIMIZATION IMAGES
# ============================================================================

def generate_power_vs_resistance_curves():
    """Image 9: Graph of power delivered vs resistance"""
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 10), sharex=True)

    # Parameters
    f = 200e3  # 200 kHz
    omega = 2 * np.pi * f
    C_mut = 8e-12  # 8 pF
    C_sh = 4e-12   # 4 pF
    C_total = C_mut + C_sh
    V_top = 350e3  # 350 kV

    # Calculate optimal resistances
    R_opt_power = 1 / (omega * C_total)
    R_opt_phase = 1 / (omega * np.sqrt(C_mut * (C_mut + C_sh)))

    # Resistance range
    R = np.logspace(3, 8, 500)  # 1 kΩ to 100 MΩ

    # Calculate power (simplified - assumes fixed V_top)
    # P ≈ 0.5 * V² * Re{Y}
    G = 1 / R
    B1 = omega * C_mut
    B2 = omega * C_sh

    # Re{Y} = G * B2² / (G² + (B1 + B2)²)
    Re_Y = G * B2**2 / (G**2 + (B1 + B2)**2)
    P = 0.5 * V_top**2 * Re_Y / 1000  # Convert to kW

    # Calculate phase angle
    Im_Y = B2 * (G**2 + B1*(B1 + B2)) / (G**2 + (B1 + B2)**2)
    phi_Z = -np.degrees(np.arctan(Im_Y / Re_Y))

    # Plot 1: Power vs Resistance
    ax1.semilogx(R/1000, P, linewidth=3, color=STYLE['primary_color'], label='Power Delivered')
    ax1.axvline(x=R_opt_power/1000, color='red', linestyle='--', linewidth=2,
               label=f'R_opt_power = {R_opt_power/1000:.1f} kΩ')
    ax1.axvline(x=R_opt_phase/1000, color='green', linestyle='--', linewidth=2,
               label=f'R_opt_phase = {R_opt_phase/1000:.1f} kΩ')

    # Mark maximum
    max_idx = np.argmax(P)
    ax1.plot(R[max_idx]/1000, P[max_idx], 'ro', markersize=12, zorder=10)
    ax1.annotate(f'Maximum Power\n{P[max_idx]:.1f} kW',
                xy=(R[max_idx]/1000, P[max_idx]),
                xytext=(R[max_idx]/5000, P[max_idx]*0.8),
                fontsize=11, fontweight='bold',
                arrowprops=dict(arrowstyle='->', color='red', lw=2),
                bbox=dict(boxstyle='round', facecolor='yellow', alpha=0.8))

    ax1.set_ylabel('Power Delivered (kW)', fontsize=STYLE['label_size'])
    ax1.set_title('Power Transfer vs Spark Resistance',
                 fontsize=STYLE['title_size'], fontweight='bold')
    ax1.grid(True, alpha=STYLE['grid_alpha'], which='both')
    ax1.legend(loc='upper right', fontsize=STYLE['legend_size'])
    set_style(ax1)

    # Plot 2: Phase angle vs Resistance
    ax2.semilogx(R/1000, phi_Z, linewidth=3, color=STYLE['secondary_color'], label='Impedance Phase φ_Z')
    ax2.axvline(x=R_opt_power/1000, color='red', linestyle='--', linewidth=2, alpha=0.7)
    ax2.axvline(x=R_opt_phase/1000, color='green', linestyle='--', linewidth=2, alpha=0.7)
    ax2.axhline(y=-45, color='orange', linestyle=':', linewidth=2, label='-45° Reference')

    # Find minimum phase magnitude
    min_phase_idx = np.argmax(phi_Z)  # Closest to 0 (least negative)
    ax2.plot(R[min_phase_idx]/1000, phi_Z[min_phase_idx], 'gs', markersize=10, zorder=10)
    ax2.annotate(f'Minimum |φ_Z|\n{phi_Z[min_phase_idx]:.1f}°',
                xy=(R[min_phase_idx]/1000, phi_Z[min_phase_idx]),
                xytext=(R[min_phase_idx]*2/1000, phi_Z[min_phase_idx] + 5),
                fontsize=10, fontweight='bold',
                arrowprops=dict(arrowstyle='->', color='green', lw=2))

    # Shade typical range
    ax2.axhspan(-75, -55, alpha=0.2, color='lightcoral', label='Typical Spark Range')

    ax2.set_xlabel('Spark Resistance R (kΩ)', fontsize=STYLE['label_size'])
    ax2.set_ylabel('Impedance Phase φ_Z (degrees)', fontsize=STYLE['label_size'])
    ax2.set_title('Impedance Phase vs Spark Resistance',
                 fontsize=STYLE['title_size'], fontweight='bold')
    ax2.grid(True, alpha=STYLE['grid_alpha'], which='both')
    ax2.legend(loc='lower right', fontsize=STYLE['legend_size'])
    set_style(ax2)

    fig.suptitle(f'Power and Phase Optimization (f = {f/1000:.0f} kHz, C_total = {C_total*1e12:.0f} pF)',
                fontsize=STYLE['title_size']+2, fontweight='bold', y=0.995)

    save_figure(fig, 'power-vs-resistance-curves.png', 'optimization')


def generate_frequency_shift_with_loading():
    """Image 13: Graph showing resonant frequency shift as spark grows"""
    fig, ax = plt.subplots(figsize=(10, 7))

    # Spark length
    L_spark = np.linspace(0, 3, 100)  # 0 to 3 meters

    # C_sh increases with length (~6.6 pF/m = 2 pF/foot)
    C_sh_0 = 3e-12  # Base capacitance (3 pF)
    C_sh_per_meter = 6.6e-12  # pF per meter
    C_sh = C_sh_0 + C_sh_per_meter * L_spark

    # Base parameters
    L_secondary = 50e-3  # 50 mH
    C_topload = 15e-12   # 15 pF
    k = 0.15  # Coupling coefficient

    # Unloaded resonance
    f0 = 1 / (2 * np.pi * np.sqrt(L_secondary * C_topload)) / 1000  # kHz

    # Loaded resonance (simplified - just secondary with increased capacitance)
    C_loaded = C_topload + C_sh
    f_loaded = 1 / (2 * np.pi * np.sqrt(L_secondary * C_loaded)) / 1000  # kHz

    # Coupled system poles (simplified)
    # Lower pole shifts down more, upper pole shifts up slightly
    f_lower = f_loaded * (1 - k/2)
    f_upper = f_loaded * (1 + k/2)

    # Plot
    ax.plot(L_spark, f_loaded, linewidth=3, color=STYLE['primary_color'], label='Loaded Resonance (secondary)')
    ax.plot(L_spark, f_lower, linewidth=2.5, color=STYLE['secondary_color'], label='Lower Pole (coupled system)')
    ax.plot(L_spark, f_upper, linewidth=2.5, color=STYLE['accent_color'], label='Upper Pole (coupled system)')
    ax.axhline(y=f0, color='black', linestyle='--', linewidth=2, label=f'Unloaded f₀ = {f0:.1f} kHz')

    # Annotations
    ax.annotate('',
                xy=(2.5, f_loaded[-1]), xytext=(2.5, f0),
                arrowprops=dict(arrowstyle='<->', color='red', lw=2))
    ax.text(2.6, (f0 + f_loaded[-1])/2,
           f'Shift:\n{f0 - f_loaded[-1]:.1f} kHz\n({(f0 - f_loaded[-1])/f0*100:.1f}%)',
           fontsize=10, fontweight='bold', color='red')

    # Mark typical operating point
    L_typical = 2.0
    idx_typical = np.argmin(np.abs(L_spark - L_typical))
    ax.plot(L_typical, f_lower[idx_typical], 'rs', markersize=12)
    ax.annotate('Typical operating\npoint (2 m spark)',
                xy=(L_typical, f_lower[idx_typical]),
                xytext=(L_typical - 0.8, f_lower[idx_typical] + 5),
                fontsize=10,
                arrowprops=dict(arrowstyle='->', color='black', lw=1.5),
                bbox=dict(boxstyle='round', facecolor='yellow', alpha=0.7))

    # Add C_sh growth annotation
    ax2 = ax.twinx()
    ax2.plot(L_spark, C_sh * 1e12, 'g--', linewidth=2, alpha=0.6, label='C_sh (pF)')
    ax2.set_ylabel('C_sh (pF)', fontsize=STYLE['label_size'], color='green')
    ax2.tick_params(axis='y', labelcolor='green')

    ax.set_xlabel('Spark Length (meters)', fontsize=STYLE['label_size'])
    ax.set_ylabel('Frequency (kHz)', fontsize=STYLE['label_size'])
    ax.set_title('Resonant Frequency Shift with Spark Loading\n(C_sh ≈ 6.6 pF/m = 2 pF/foot)',
                fontsize=STYLE['title_size'], fontweight='bold')
    ax.grid(True, alpha=STYLE['grid_alpha'])
    ax.legend(loc='upper right', fontsize=STYLE['legend_size'])

    # Key insight box
    textstr = 'Critical: For accurate power\nmeasurement, retune to loaded\npole for each spark length!'
    props = dict(boxstyle='round', facecolor='wheat', alpha=0.9, edgecolor='red', linewidth=2)
    ax.text(0.15, f0 - 3, textstr, fontsize=11, verticalalignment='top', bbox=props, fontweight='bold')

    save_figure(fig, 'frequency-shift-with-loading.png', 'optimization')


# ============================================================================
# PART 3: SPARK PHYSICS IMAGES
# ============================================================================

def generate_energy_budget_breakdown():
    """Image 18: Pie chart showing energy distribution per meter"""
    fig, ax = plt.subplots(figsize=(9, 9))

    # Energy components (percentage)
    labels = [
        'Ionization Energy\n(40-50%)',
        'Channel Heating\n(20-30%)',
        'Radiation Losses\n(10-20%)',
        'Shock Wave /\nAcoustic (5-10%)',
        'Electrohydrodynamic\nWork (5-10%)'
    ]
    sizes = [45, 25, 15, 8, 7]
    colors = ['#FF6B6B', '#4ECDC4', '#45B7D1', '#FFA07A', '#98D8C8']
    explode = (0.05, 0, 0, 0, 0)  # Emphasize ionization

    wedges, texts, autotexts = ax.pie(sizes, explode=explode, labels=labels, colors=colors,
                                       autopct='%1.1f%%', startangle=90, textprops={'fontsize': 11})

    # Enhance text
    for autotext in autotexts:
        autotext.set_color('white')
        autotext.set_fontweight('bold')
        autotext.set_fontsize(12)

    for text in texts:
        text.set_fontsize(11)
        text.set_fontweight('bold')

    ax.set_title('Energy Distribution per Meter of Spark Growth\n(QCW Mode, ε ≈ 10 J/m)',
                fontsize=STYLE['title_size']+1, fontweight='bold', pad=20)

    # Add note
    note = ('Theoretical minimum: ~0.5 J/m\n'
            'Actual QCW: 5-15 J/m\n'
            'Burst mode: 30-100+ J/m')
    ax.text(1.4, -1.1, note, fontsize=10,
           bbox=dict(boxstyle='round', facecolor='lightyellow', alpha=0.8),
           verticalalignment='top')

    ax.axis('equal')

    save_figure(fig, 'energy-budget-breakdown.png', 'spark-physics')


def generate_thermal_diffusion_vs_diameter():
    """Image 20: Graph of thermal time constant vs channel diameter"""
    fig, ax = plt.subplots(figsize=(10, 7))

    # Channel diameter range
    d = np.logspace(-5, -2, 300)  # 10 μm to 10 mm

    # Thermal diffusivity of air
    alpha_thermal = 2e-5  # m²/s

    # Thermal time constant: τ = d² / (4α)
    tau = d**2 / (4 * alpha_thermal) * 1000  # Convert to ms

    # Plot
    ax.loglog(d * 1e6, tau, linewidth=3, color=STYLE['primary_color'], label='τ = d² / (4α)')

    # Mark key points
    points = [
        (50e-6, 'Thin streamer\n(50 μm)', 'top'),
        (100e-6, 'Typical streamer\n(100 μm)', 'top'),
        (1e-3, 'Thin leader\n(1 mm)', 'bottom'),
        (5e-3, 'Thick leader\n(5 mm)', 'bottom'),
    ]

    for d_point, label, valign in points:
        tau_point = d_point**2 / (4 * alpha_thermal) * 1000
        ax.plot(d_point * 1e6, tau_point, 'ro', markersize=10)

        y_offset = tau_point * 2.5 if valign == 'top' else tau_point / 2.5
        ax.annotate(f'{label}\nτ ≈ {tau_point:.2f} ms',
                   xy=(d_point * 1e6, tau_point),
                   xytext=(d_point * 1e6, y_offset),
                   fontsize=9, ha='center',
                   arrowprops=dict(arrowstyle='->', color='black', lw=1),
                   bbox=dict(boxstyle='round', facecolor='yellow', alpha=0.7))

    # Shade regimes
    ax.axvspan(10, 200, alpha=0.2, color='purple', label='Streamer Regime')
    ax.axvspan(500, 10000, alpha=0.2, color='orange', label='Leader Regime')

    # Add convection note
    textstr = ('Note: Observed persistence\nlonger than pure thermal\ndiffusion due to:\n'
               '• Convection\n'
               '• Ionization memory\n'
               '• Buoyancy effects')
    props = dict(boxstyle='round', facecolor='lightblue', alpha=0.8)
    ax.text(30, 200, textstr, fontsize=10, verticalalignment='top', bbox=props)

    ax.set_xlabel('Channel Diameter d (μm)', fontsize=STYLE['label_size'])
    ax.set_ylabel('Thermal Time Constant τ (ms)', fontsize=STYLE['label_size'])
    ax.set_title('Thermal Diffusion Time vs Channel Diameter\n(α = 2×10⁻⁵ m²/s for air)',
                fontsize=STYLE['title_size'], fontweight='bold')
    ax.grid(True, alpha=STYLE['grid_alpha'], which='both')
    ax.legend(loc='upper left', fontsize=STYLE['legend_size'])

    save_figure(fig, 'thermal-diffusion-vs-diameter.png', 'spark-physics')


def generate_voltage_division_vs_length_plot():
    """Image 24: Graph showing how V_tip decreases as spark grows"""
    fig, ax = plt.subplots(figsize=(10, 7))

    # Parameters
    C_mut = 10e-12  # 10 pF
    C_sh_per_meter = 6.6e-12  # 6.6 pF/m
    V_topload = 350  # kV
    E_propagation = 0.5  # MV/m

    # Spark length
    L = np.linspace(0, 3, 300)  # 0 to 3 meters

    # C_sh grows with length
    C_sh = C_sh_per_meter * L
    C_sh[0] = 0.1e-12  # Avoid division by zero

    # Voltage division (open-circuit limit)
    V_tip_ratio = C_mut / (C_mut + C_sh)
    V_tip = V_topload * V_tip_ratio

    # Electric field at tip (simplified: E_tip ≈ V_tip / L)
    E_tip = V_tip / L
    E_tip[0] = V_topload / 0.001  # Avoid infinity at L=0

    # Plot V_tip ratio
    ax1 = ax
    ax1.plot(L, V_tip_ratio, linewidth=3, color=STYLE['primary_color'], label='V_tip / V_topload')
    ax1.set_xlabel('Spark Length L (meters)', fontsize=STYLE['label_size'])
    ax1.set_ylabel('Voltage Ratio V_tip / V_topload', fontsize=STYLE['label_size'], color=STYLE['primary_color'])
    ax1.tick_params(axis='y', labelcolor=STYLE['primary_color'])
    ax1.grid(True, alpha=STYLE['grid_alpha'])

    # Add E_tip on secondary axis
    ax2 = ax1.twinx()
    ax2.plot(L, E_tip, linewidth=2.5, color=STYLE['secondary_color'], linestyle='--', label='E_tip (MV/m)')
    ax2.axhline(y=E_propagation, color='green', linestyle=':', linewidth=2, label=f'E_propagation = {E_propagation} MV/m')
    ax2.set_ylabel('Electric Field E_tip (MV/m)', fontsize=STYLE['label_size'], color=STYLE['secondary_color'])
    ax2.tick_params(axis='y', labelcolor=STYLE['secondary_color'])
    ax2.set_ylim(0, 5)

    # Find where E_tip = E_propagation (growth stalls)
    idx_stall = np.argmin(np.abs(E_tip - E_propagation))
    L_stall = L[idx_stall]

    ax1.axvline(x=L_stall, color='red', linestyle='--', linewidth=2, alpha=0.7)
    ax1.annotate(f'Growth Stalls\n(E_tip = E_prop)\nL ≈ {L_stall:.2f} m',
                xy=(L_stall, V_tip_ratio[idx_stall]),
                xytext=(L_stall + 0.5, V_tip_ratio[idx_stall] + 0.2),
                fontsize=11, fontweight='bold', color='red',
                arrowprops=dict(arrowstyle='->', color='red', lw=2),
                bbox=dict(boxstyle='round', facecolor='yellow', alpha=0.8))

    # Annotations
    ax1.annotate('Sub-linear scaling:\nV_tip drops even if\nV_topload maintained',
                xy=(1.5, V_tip_ratio[150]),
                xytext=(2.0, 0.7),
                fontsize=10,
                arrowprops=dict(arrowstyle='->', color='black', lw=1.5),
                bbox=dict(boxstyle='round', facecolor='lightblue', alpha=0.7))

    # Formula
    formula = f'V_tip = V_topload × C_mut/(C_mut + C_sh)\nC_sh ≈ {C_sh_per_meter*1e12:.1f} pF/m'
    ax1.text(0.1, 0.2, formula, fontsize=10,
            bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8))

    ax1.set_title('Capacitive Divider Effect: V_tip vs Spark Length\n(Open-circuit limit, R → ∞)',
                 fontsize=STYLE['title_size'], fontweight='bold')

    # Combine legends
    lines1, labels1 = ax1.get_legend_handles_labels()
    lines2, labels2 = ax2.get_legend_handles_labels()
    ax1.legend(lines1 + lines2, labels1 + labels2, loc='upper right', fontsize=STYLE['legend_size'])

    save_figure(fig, 'voltage-division-vs-length-plot.png', 'spark-physics')


def generate_length_vs_energy_scaling():
    """Image 26: Log-log plot showing L vs E scaling for different modes"""
    fig, ax = plt.subplots(figsize=(10, 8))

    # Energy range
    E = np.logspace(0, 3, 100)  # 1 J to 1000 J

    # Different scaling relationships
    # Burst mode: L ∝ √E (slope = 0.5)
    epsilon_burst = 50  # J/m
    L_burst = np.sqrt(E / epsilon_burst)

    # QCW: L ∝ E^0.7 (slope = 0.7)
    epsilon_qcw = 10  # J/m
    L_qcw = (E / epsilon_qcw)**0.7 / 3  # Normalize

    # Ideal linear: L ∝ E (slope = 1.0)
    epsilon_ideal = 5  # J/m
    L_ideal = E / epsilon_ideal / 10  # Normalize

    # Plot
    ax.loglog(E, L_burst, linewidth=3, color=STYLE['secondary_color'],
             label='Burst Mode: L ∝ √E (slope = 0.5)', marker='o', markevery=10)
    ax.loglog(E, L_qcw, linewidth=3, color=STYLE['primary_color'],
             label='QCW: L ∝ E^0.7 (slope = 0.7)', marker='s', markevery=10)
    ax.loglog(E, L_ideal, linewidth=2, color='green', linestyle='--',
             label='Ideal: L ∝ E (slope = 1.0)', marker='^', markevery=10)

    # Add slope annotations
    for E_point, L_point, slope, color, offset in [
        (10, L_burst[np.argmin(np.abs(E - 10))], '0.5', STYLE['secondary_color'], (1.5, 0.7)),
        (50, L_qcw[np.argmin(np.abs(E - 50))], '0.7', STYLE['primary_color'], (1.5, 1.0)),
        (200, L_ideal[np.argmin(np.abs(E - 200))], '1.0', 'green', (1.3, 1.0)),
    ]:
        ax.annotate(f'Slope = {slope}',
                   xy=(E_point, L_point),
                   xytext=(E_point * offset[0], L_point * offset[1]),
                   fontsize=10, color=color, fontweight='bold',
                   arrowprops=dict(arrowstyle='->', color=color, lw=1.5))

    # Add data points (simulated realistic observations)
    # Burst mode data
    E_burst_data = np.array([5, 10, 20, 50, 100])
    L_burst_data = np.sqrt(E_burst_data / 55) * (1 + 0.1 * np.random.randn(5))
    ax.plot(E_burst_data, L_burst_data, 'r*', markersize=12, label='Burst Mode (measured)')

    # QCW data
    E_qcw_data = np.array([10, 25, 50, 100, 200])
    L_qcw_data = (E_qcw_data / 10)**0.72 / 3 * (1 + 0.08 * np.random.randn(5))
    ax.plot(E_qcw_data, L_qcw_data, 'b*', markersize=12, label='QCW (measured)')

    ax.set_xlabel('Energy E (Joules)', fontsize=STYLE['label_size'])
    ax.set_ylabel('Spark Length L (meters)', fontsize=STYLE['label_size'])
    ax.set_title('Freau\'s Empirical Scaling: Spark Length vs Energy\n(Sub-linear scaling due to capacitive divider)',
                fontsize=STYLE['title_size'], fontweight='bold')
    ax.grid(True, alpha=STYLE['grid_alpha'], which='both')
    ax.legend(loc='upper left', fontsize=STYLE['legend_size'])

    # Physical explanation box
    textstr = ('Physical Explanation:\n'
               '• Burst: Voltage-limited\n'
               '  (V_tip drops with length)\n'
               '• QCW: Better voltage\n'
               '  maintenance via ramping\n'
               '• Ideal: Constant ε,\n'
               '  constant E_tip')
    props = dict(boxstyle='round', facecolor='lightyellow', alpha=0.9)
    ax.text(2, 3, textstr, fontsize=10, verticalalignment='top', bbox=props)

    save_figure(fig, 'length-vs-energy-scaling.png', 'spark-physics')


# ============================================================================
# PART 4: ADVANCED MODELING IMAGES
# ============================================================================

def generate_capacitance_matrix_heatmap():
    """Image 34: Heatmap visualization of 11×11 capacitance matrix"""
    fig, ax = plt.subplots(figsize=(10, 9))

    # Create realistic 11×11 capacitance matrix (topload + 10 segments)
    n = 11
    C_matrix = np.zeros((n, n))

    # Diagonal elements (large positive)
    for i in range(n):
        if i == 0:  # Topload
            C_matrix[i, i] = 25.0
        else:  # Segments (decreasing toward tip)
            C_matrix[i, i] = 15.0 - i * 0.8

    # Off-diagonal elements (negative, stronger for adjacent)
    for i in range(n):
        for j in range(i+1, n):
            distance = abs(i - j)
            if distance == 1:  # Adjacent
                C_matrix[i, j] = -3.5 + np.random.rand() * 0.5
            elif distance == 2:
                C_matrix[i, j] = -1.2 + np.random.rand() * 0.3
            elif distance == 3:
                C_matrix[i, j] = -0.5 + np.random.rand() * 0.2
            else:
                C_matrix[i, j] = -0.1 - np.random.rand() * 0.05

            C_matrix[j, i] = C_matrix[i, j]  # Symmetric

    # Plot heatmap
    im = ax.imshow(C_matrix, cmap='RdBu_r', aspect='equal')

    # Colorbar
    cbar = plt.colorbar(im, ax=ax)
    cbar.set_label('Capacitance (pF)', fontsize=STYLE['label_size'])

    # Labels
    labels = ['Top'] + [f'Seg{i}' for i in range(1, n)]
    ax.set_xticks(np.arange(n))
    ax.set_yticks(np.arange(n))
    ax.set_xticklabels(labels)
    ax.set_yticklabels(labels)

    # Rotate x labels
    plt.setp(ax.get_xticklabels(), rotation=45, ha="right", rotation_mode="anchor")

    # Add grid
    ax.set_xticks(np.arange(n)-.5, minor=True)
    ax.set_yticks(np.arange(n)-.5, minor=True)
    ax.grid(which="minor", color="gray", linestyle='-', linewidth=0.5)

    # Annotations
    ax.text(-1.5, n/2, 'Rows', fontsize=12, rotation=90, ha='center', va='center', fontweight='bold')
    ax.text(n/2, -1.5, 'Columns', fontsize=12, ha='center', va='center', fontweight='bold')

    # Add some value annotations
    for i in range(min(3, n)):
        for j in range(min(3, n)):
            text = ax.text(j, i, f'{C_matrix[i, j]:.1f}',
                          ha="center", va="center", color="black", fontsize=8)

    ax.set_title('Maxwell Capacitance Matrix (11×11)\nTopload + 10 Spark Segments',
                fontsize=STYLE['title_size'], fontweight='bold')

    # Add notes
    note = ('• Diagonal: Large positive (self-capacitance)\n'
            '• Off-diagonal: Negative (mutual)\n'
            '• Adjacent elements: Stronger coupling\n'
            '• Matrix is symmetric')
    fig.text(0.02, 0.02, note, fontsize=9,
            bbox=dict(boxstyle='round', facecolor='lightyellow', alpha=0.8),
            verticalalignment='bottom')

    save_figure(fig, 'capacitance-matrix-heatmap.png', 'advanced-modeling')


def generate_resistance_taper_initialization():
    """Image 36: Graph showing initial resistance distribution"""
    fig, ax = plt.subplots(figsize=(10, 7))

    # Position along spark
    position = np.linspace(0, 1, 100)  # 0 = base, 1 = tip

    # Three initialization strategies
    R_base = 10e3  # 10 kΩ
    R_tip = 1e6   # 1 MΩ

    # 1. Uniform (wrong)
    R_uniform = np.ones_like(position) * np.sqrt(R_base * R_tip)

    # 2. Linear taper
    R_linear = R_base + (R_tip - R_base) * position

    # 3. Quadratic taper (recommended)
    R_quadratic = R_base + (R_tip - R_base) * position**2

    # Physical bounds
    R_min = R_base + (10e3 - R_base) * position
    R_max = 100e3 + (100e6 - 100e3) * position**2

    # Plot
    ax.semilogy(position, R_uniform, linewidth=2, linestyle=':', color='red',
               label='Uniform (poor)', alpha=0.7)
    ax.semilogy(position, R_linear, linewidth=2.5, linestyle='--', color=STYLE['primary_color'],
               label='Linear Taper (better)')
    ax.semilogy(position, R_quadratic, linewidth=3, color=STYLE['accent_color'],
               label='Quadratic Taper (recommended)')

    # Shade physical bounds
    ax.fill_between(position, R_min, R_max, alpha=0.15, color='gray', label='Physical Bounds')
    ax.semilogy(position, R_min, 'k--', linewidth=1, alpha=0.5)
    ax.semilogy(position, R_max, 'k--', linewidth=1, alpha=0.5)

    # Annotations
    ax.annotate('Hot leader\n(low R)', xy=(0, R_base), xytext=(0.1, R_base/3),
               fontsize=10, arrowprops=dict(arrowstyle='->', color='black', lw=1))
    ax.annotate('Cold streamer\n(high R)', xy=(1, R_tip), xytext=(0.85, R_tip*3),
               fontsize=10, arrowprops=dict(arrowstyle='->', color='black', lw=1))

    # Formula
    formula = ('Quadratic: R[i] = R_base + (R_tip - R_base) × pos²\n'
               f'R_base = {R_base/1000:.0f} kΩ, R_tip = {R_tip/1e6:.1f} MΩ')
    ax.text(0.5, 2e6, formula, fontsize=10, ha='center',
           bbox=dict(boxstyle='round', facecolor='yellow', alpha=0.7))

    ax.set_xlabel('Normalized Position (0 = base, 1 = tip)', fontsize=STYLE['label_size'])
    ax.set_ylabel('Resistance R[i] (Ω)', fontsize=STYLE['label_size'])
    ax.set_title('Resistance Distribution Initialization Strategies',
                fontsize=STYLE['title_size'], fontweight='bold')
    ax.grid(True, alpha=STYLE['grid_alpha'], which='both')
    ax.legend(loc='upper left', fontsize=STYLE['legend_size'])

    save_figure(fig, 'resistance-taper-initialization.png', 'advanced-modeling')


def generate_power_distribution_along_spark():
    """Image 38: Bar chart showing power dissipation per segment"""
    fig, ax = plt.subplots(figsize=(12, 7))

    # Segment numbers
    segments = np.arange(1, 11)

    # Realistic power distribution (higher at base, peaks at segment 2-3, decays to tip)
    power = np.array([280, 380, 340, 280, 220, 160, 110, 70, 40, 20])  # kW

    # Calculate cumulative
    power_cumulative = np.cumsum(power)
    total_power = power_cumulative[-1]

    # Bar chart
    bars = ax.bar(segments, power, color=STYLE['primary_color'], edgecolor='black', linewidth=1.5)

    # Color gradient (hot at base, cool at tip)
    colors = plt.cm.hot(np.linspace(0.8, 0.2, len(segments)))
    for bar, color in zip(bars, colors):
        bar.set_facecolor(color)

    # Add percentage labels on bars
    for i, (seg, p) in enumerate(zip(segments, power)):
        percentage = p / total_power * 100
        ax.text(seg, p + 20, f'{percentage:.1f}%', ha='center', fontsize=9, fontweight='bold')

    # Cumulative line
    ax2 = ax.twinx()
    ax2.plot(segments, power_cumulative / total_power * 100, 'bo-', linewidth=2.5, markersize=8,
            label='Cumulative (%)')
    ax2.set_ylabel('Cumulative Power (%)', fontsize=STYLE['label_size'], color='blue')
    ax2.tick_params(axis='y', labelcolor='blue')
    ax2.set_ylim(0, 110)
    ax2.axhline(y=100, color='blue', linestyle='--', alpha=0.5)

    # Annotations
    ax.annotate('Peak power\nin segments 2-3', xy=(2.5, 380), xytext=(4, 450),
               fontsize=11, fontweight='bold',
               arrowprops=dict(arrowstyle='->', color='red', lw=2),
               bbox=dict(boxstyle='round', facecolor='yellow', alpha=0.8))

    ax.annotate('Base segments:\n68% of total power', xy=(1, 280), xytext=(0.3, 350),
               fontsize=10,
               arrowprops=dict(arrowstyle='->', color='black', lw=1.5))

    ax.annotate('Tip: Only 5%\nof total power', xy=(10, 20), xytext=(8.5, 100),
               fontsize=10,
               arrowprops=dict(arrowstyle='->', color='black', lw=1.5))

    ax.set_xlabel('Segment Number (1 = base, 10 = tip)', fontsize=STYLE['label_size'])
    ax.set_ylabel('Power Dissipated (kW)', fontsize=STYLE['label_size'])
    ax.set_title(f'Power Distribution Along Spark (Total = {total_power:.0f} kW)',
                fontsize=STYLE['title_size'], fontweight='bold')
    ax.set_xticks(segments)
    ax.grid(True, alpha=STYLE['grid_alpha'], axis='y')

    # Total power box
    textstr = f'Total Power: {total_power:.0f} kW\nBase (1-3): 68%\nMiddle (4-7): 27%\nTip (8-10): 5%'
    props = dict(boxstyle='round', facecolor='lightblue', alpha=0.9)
    ax.text(8.5, 400, textstr, fontsize=11, bbox=props, fontweight='bold')

    save_figure(fig, 'power-distribution-along-spark.png', 'advanced-modeling')


def generate_current_attenuation_plot():
    """Image 39: Graph of current magnitude along spark"""
    fig, ax = plt.subplots(figsize=(10, 7))

    # Position along spark
    position = np.linspace(0, 2.5, 100)  # 0 to 2.5 meters

    # Current attenuation (exponential-like decay)
    # More current flows through capacitances to ground as we move along
    I_normalized = np.exp(-0.6 * position)

    # Add some realistic variation
    I_normalized = I_normalized * (1 + 0.03 * np.sin(20 * position))

    # Plot
    ax.plot(position, I_normalized * 100, linewidth=3, color=STYLE['primary_color'])
    ax.fill_between(position, 0, I_normalized * 100, alpha=0.3, color=STYLE['primary_color'])

    # Mark segment boundaries (10 segments)
    segment_positions = np.linspace(0, 2.5, 11)
    for i, pos in enumerate(segment_positions):
        I_val = np.exp(-0.6 * pos) * 100
        ax.plot([pos, pos], [0, I_val], 'k--', linewidth=1, alpha=0.4)
        if i < len(segment_positions) - 1:
            ax.text(pos, -8, f'{i+1}', ha='center', fontsize=9, fontweight='bold')

    # Mark key points
    key_points = [
        (0, 'Base: 100%'),
        (0.83, 'Middle: 60%'),
        (1.67, '3/4 point: 36%'),
        (2.5, 'Tip: 22%'),
    ]

    for pos, label in key_points:
        I_val = np.exp(-0.6 * pos) * 100
        ax.plot(pos, I_val, 'ro', markersize=10)
        ax.annotate(label, xy=(pos, I_val), xytext=(pos, I_val + 15),
                   fontsize=10, ha='center', fontweight='bold',
                   arrowprops=dict(arrowstyle='->', color='red', lw=1.5))

    ax.set_xlabel('Position Along Spark (meters)', fontsize=STYLE['label_size'])
    ax.set_ylabel('Normalized Current |I| / |I_base| (%)', fontsize=STYLE['label_size'])
    ax.set_title('Current Attenuation Along Distributed Spark Model\n(Current diverted to ground through C_sh)',
                fontsize=STYLE['title_size'], fontweight='bold')
    ax.set_xlim(0, 2.5)
    ax.set_ylim(0, 110)
    ax.grid(True, alpha=STYLE['grid_alpha'])

    # Segment labels
    ax.text(1.25, -15, 'Segment Number', ha='center', fontsize=11, fontweight='bold')

    # Physical explanation
    textstr = ('Current decreases due to:\n'
               '• Displacement current through C_sh\n'
               '• Distributed capacitive shunting\n'
               '• Not ohmic attenuation (R is small)')
    props = dict(boxstyle='round', facecolor='lightyellow', alpha=0.9)
    ax.text(1.8, 75, textstr, fontsize=10, bbox=props)

    save_figure(fig, 'current-attenuation-plot.png', 'advanced-modeling')


# ============================================================================
# SHARED IMAGES
# ============================================================================

def generate_complex_number_review():
    """Image 45: Quick reference for complex number operations"""
    fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(14, 12))

    # Example complex number
    a = 3
    b = 4
    z = complex(a, b)
    r = abs(z)
    theta = np.angle(z)

    # Quadrant 1: Rectangular form
    ax1.arrow(0, 0, a, b, head_width=0.3, head_length=0.2, fc='black', ec='black', linewidth=2)
    ax1.plot([0, a], [0, 0], 'b--', linewidth=2, label='Real part (a)')
    ax1.plot([a, a], [0, b], 'r--', linewidth=2, label='Imaginary part (b)')
    ax1.text(a/2, -0.5, f'a = {a}', fontsize=12, ha='center', color='blue', fontweight='bold')
    ax1.text(a + 0.3, b/2, f'b = {b}', fontsize=12, ha='left', color='red', fontweight='bold')
    ax1.text(a/2 + 0.3, b/2 + 0.5, f'z = a + jb\n= {a} + j{b}', fontsize=13, fontweight='bold',
            bbox=dict(boxstyle='round', facecolor='yellow', alpha=0.7))
    ax1.set_xlim(-1, 6)
    ax1.set_ylim(-1, 6)
    ax1.set_xlabel('Real Axis', fontsize=11)
    ax1.set_ylabel('Imaginary Axis', fontsize=11)
    ax1.set_title('Rectangular Form: z = a + jb', fontsize=STYLE['title_size'], fontweight='bold')
    ax1.axhline(y=0, color='k', linewidth=0.8)
    ax1.axvline(x=0, color='k', linewidth=0.8)
    ax1.grid(True, alpha=0.3)
    ax1.legend(loc='upper left')
    ax1.set_aspect('equal')

    # Quadrant 2: Polar form
    ax2.arrow(0, 0, a, b, head_width=0.3, head_length=0.2, fc='black', ec='black', linewidth=2)
    arc = Wedge((0, 0), 1.5, 0, np.degrees(theta), facecolor='lightgreen', alpha=0.5, edgecolor='black')
    ax2.add_patch(arc)
    ax2.plot([0, r], [0, 0], 'g--', linewidth=1, alpha=0.5)
    ax2.text(r/2, -0.5, f'r = |z| = {r:.2f}', fontsize=11, ha='center', color='green', fontweight='bold')
    ax2.text(0.8, 0.6, f'θ = {np.degrees(theta):.1f}°', fontsize=11, color='green', fontweight='bold')
    ax2.text(a/2 + 0.3, b/2 + 0.8, f'z = r∠θ\n= {r:.2f}∠{np.degrees(theta):.1f}°', fontsize=13, fontweight='bold',
            bbox=dict(boxstyle='round', facecolor='lightblue', alpha=0.7))
    ax2.set_xlim(-1, 6)
    ax2.set_ylim(-1, 6)
    ax2.set_xlabel('Real Axis', fontsize=11)
    ax2.set_ylabel('Imaginary Axis', fontsize=11)
    ax2.set_title('Polar Form: z = r∠θ', fontsize=STYLE['title_size'], fontweight='bold')
    ax2.axhline(y=0, color='k', linewidth=0.8)
    ax2.axvline(x=0, color='k', linewidth=0.8)
    ax2.grid(True, alpha=0.3)
    ax2.set_aspect('equal')

    # Quadrant 3: Conversions
    ax3.axis('off')
    conversion_text = f"""
CONVERSIONS

Rectangular → Polar:
  r = √(a² + b²) = {r:.2f}
  θ = atan(b/a) = {np.degrees(theta):.1f}°

Polar → Rectangular:
  a = r cos(θ) = {r:.2f} × cos({np.degrees(theta):.1f}°) = {a:.2f}
  b = r sin(θ) = {r:.2f} × sin({np.degrees(theta):.1f}°) = {b:.2f}

Euler Form:
  z = r e^(jθ) = {r:.2f} e^(j{theta:.3f})

Complex Conjugate:
  z* = a - jb = {a} - j{b}
  z* = r∠(-θ) = {r:.2f}∠{-np.degrees(theta):.1f}°
"""
    ax3.text(0.1, 0.9, conversion_text, fontsize=11, fontfamily='monospace',
            verticalalignment='top',
            bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.9))

    # Quadrant 4: Operations
    ax4.axis('off')
    operations_text = """
OPERATIONS

Addition/Subtraction (use rectangular):
  z₁ + z₂ = (a₁ + a₂) + j(b₁ + b₂)

Multiplication (use polar):
  z₁ × z₂ = r₁r₂ ∠(θ₁ + θ₂)

Division (use polar):
  z₁ / z₂ = (r₁/r₂) ∠(θ₁ - θ₂)

Magnitude:
  |z| = r = √(a² + b²)

Power Calculation (AC circuits):
  P = ½ Re{V × I*}
  (Use conjugate of current!)

Important for Tesla Coils:
  j = √(-1) (imaginary unit)
  jωL → inductive reactance (+j)
  -j/(ωC) → capacitive reactance (-j)
"""
    ax4.text(0.1, 0.9, operations_text, fontsize=10, fontfamily='monospace',
            verticalalignment='top',
            bbox=dict(boxstyle='round', facecolor='lightgreen', alpha=0.9))

    fig.suptitle('Complex Number Quick Reference for AC Circuit Analysis',
                fontsize=STYLE['title_size']+2, fontweight='bold', y=0.98)

    plt.tight_layout()
    save_figure(fig, 'complex-number-review.png', 'shared')


# ============================================================================
# ADDITIONAL IMAGES: Comparison Tables and Simple Diagrams
# ============================================================================

def generate_drsstc_operating_modes():
    """Image 14: Three timing diagrams showing different DRSSTC operating modes"""
    fig, axes = plt.subplots(3, 1, figsize=(12, 9))
    fig.suptitle('DRSSTC Operating Modes Comparison', fontsize=16, fontweight='bold')

    time_ms = np.linspace(0, 25, 1000)

    # Mode 1: Fixed frequency
    ax = axes[0]
    freq_fixed = 200  # kHz
    drive_fixed = 0.8 * np.sin(2 * np.pi * freq_fixed * time_ms / 1000)
    drive_fixed[time_ms > 20] = 0  # Pulse ends

    ax.plot(time_ms, drive_fixed, 'b-', linewidth=2, label='Drive Signal (Fixed f)')
    ax.axhline(y=0, color='k', linestyle='-', linewidth=0.5)
    ax.set_ylabel('Amplitude', fontsize=11)
    ax.set_title('(a) Fixed Frequency Mode', fontsize=12, fontweight='bold')
    ax.text(22, 0.5, 'Pro: Simple\nCon: Detuning\nwith loading', fontsize=9,
            bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8))
    ax.grid(True, alpha=0.3)
    ax.set_xlim(0, 25)
    ax.legend(loc='upper left', fontsize=9)

    # Mode 2: PLL tracking
    ax = axes[1]
    # Frequency decreases as spark grows
    freq_pll = 200 - 25 * (time_ms / 20)  # 200 -> 175 kHz
    freq_pll[time_ms > 20] = 200
    phase_pll = np.cumsum(2 * np.pi * freq_pll / 1000 * np.mean(np.diff(time_ms)))
    drive_pll = 0.8 * np.sin(phase_pll)
    drive_pll[time_ms > 20] = 0

    ax.plot(time_ms, drive_pll, 'g-', linewidth=2, label='Drive Signal (PLL Tracking)')
    ax.axhline(y=0, color='k', linestyle='-', linewidth=0.5)
    ax.set_ylabel('Amplitude', fontsize=11)
    ax.set_title('(b) PLL Tracking Mode', fontsize=12, fontweight='bold')
    ax.text(22, 0.5, 'Pro: Follows\nresonance\nCon: Complex', fontsize=9,
            bbox=dict(boxstyle='round', facecolor='lightgreen', alpha=0.8))
    ax.grid(True, alpha=0.3)
    ax.set_xlim(0, 25)
    ax.legend(loc='upper left', fontsize=9)

    # Mode 3: Programmed sweep
    ax = axes[2]
    freq_sweep = 200 - 30 * (time_ms / 20)**0.7  # Predetermined curve
    freq_sweep[time_ms > 20] = 200
    phase_sweep = np.cumsum(2 * np.pi * freq_sweep / 1000 * np.mean(np.diff(time_ms)))
    drive_sweep = 0.8 * np.sin(phase_sweep)
    drive_sweep[time_ms > 20] = 0

    ax.plot(time_ms, drive_sweep, 'r-', linewidth=2, label='Drive Signal (Programmed)')
    ax.axhline(y=0, color='k', linestyle='-', linewidth=0.5)
    ax.set_ylabel('Amplitude', fontsize=11)
    ax.set_xlabel('Time (ms)', fontsize=11)
    ax.set_title('(c) Programmed Sweep Mode', fontsize=12, fontweight='bold')
    ax.text(22, 0.5, 'Pro: Optimal\ntrajectory\nCon: Tuning', fontsize=9,
            bbox=dict(boxstyle='round', facecolor='lightcoral', alpha=0.8))
    ax.grid(True, alpha=0.3)
    ax.set_xlim(0, 25)
    ax.legend(loc='upper left', fontsize=9)

    plt.tight_layout()
    save_figure(fig, 'drsstc-operating-modes.png', 'optimization')


def generate_loaded_pole_analysis():
    """Image 15: Frequency domain showing coupled resonances"""
    fig, ax = plt.subplots(figsize=(10, 7))

    freq_khz = np.linspace(150, 250, 500)
    f0 = 200  # Unloaded resonance

    # Unloaded: sharp twin peaks
    Q_unloaded = 50
    pole1_unloaded = 190
    pole2_unloaded = 210
    response_unloaded = (100 / (1 + Q_unloaded**2 * ((freq_khz - pole1_unloaded)/pole1_unloaded)**2) +
                         100 / (1 + Q_unloaded**2 * ((freq_khz - pole2_unloaded)/pole2_unloaded)**2))

    # Loaded: broader, shifted peaks
    Q_loaded = 15
    pole1_loaded = 175
    pole2_loaded = 215
    response_loaded = (80 / (1 + Q_loaded**2 * ((freq_khz - pole1_loaded)/pole1_loaded)**2) +
                       80 / (1 + Q_loaded**2 * ((freq_khz - pole2_loaded)/pole2_loaded)**2))

    ax.plot(freq_khz, response_unloaded, 'b-', linewidth=2.5, label='Unloaded (No Spark)', alpha=0.7)
    ax.plot(freq_khz, response_loaded, 'r-', linewidth=2.5, label='Loaded (With Spark)', alpha=0.7)

    # Mark operating points
    ax.axvline(x=f0, color='blue', linestyle='--', linewidth=1.5, alpha=0.5,
               label=f'Fixed f0 = {f0} kHz (Wrong!)')
    ax.axvline(x=pole1_loaded, color='red', linestyle='--', linewidth=1.5, alpha=0.7,
               label=f'Track Loaded Pole = {pole1_loaded} kHz (Right!)')

    # Annotations
    ax.annotate('Unloaded peaks:\nSharp, symmetric', xy=(pole1_unloaded, 90),
                xytext=(160, 150), fontsize=10,
                bbox=dict(boxstyle='round', facecolor='lightblue', alpha=0.8),
                arrowprops=dict(arrowstyle='->', color='blue', lw=1.5))

    ax.annotate('Loaded peaks:\nBroader, shifted', xy=(pole1_loaded, 70),
                xytext=(160, 90), fontsize=10,
                bbox=dict(boxstyle='round', facecolor='lightcoral', alpha=0.8),
                arrowprops=dict(arrowstyle='->', color='red', lw=1.5))

    ax.set_xlabel('Frequency (kHz)', fontsize=12)
    ax.set_ylabel('|V_topload| (kV)', fontsize=12)
    ax.set_title('Loaded Pole Analysis: Frequency Tracking is Critical', fontsize=14, fontweight='bold')
    ax.legend(loc='upper right', fontsize=10)
    ax.grid(True, alpha=0.3)
    ax.set_xlim(150, 250)
    ax.set_ylim(0, 200)

    plt.tight_layout()
    save_figure(fig, 'loaded-pole-analysis.png', 'optimization')


def generate_epsilon_by_mode_comparison():
    """Image 19: Bar chart comparing epsilon values by operating mode"""
    fig, ax = plt.subplots(figsize=(10, 7))

    modes = ['QCW\n(Continuous)', 'Hybrid\nDRSSTC', 'Hard-Pulsed\nBurst']
    epsilon_mean = [10, 30, 60]
    epsilon_err_low = [5, 10, 30]
    epsilon_err_high = [5, 10, 40]
    colors = ['green', 'orange', 'red']

    x_pos = np.arange(len(modes))
    bars = ax.bar(x_pos, epsilon_mean, color=colors, alpha=0.7, edgecolor='black', linewidth=1.5)

    # Error bars
    ax.errorbar(x_pos, epsilon_mean,
                yerr=[epsilon_err_low, epsilon_err_high],
                fmt='none', color='black', capsize=10, capthick=2, linewidth=2)

    # Add value labels
    for i, (mean, color) in enumerate(zip(epsilon_mean, colors)):
        ax.text(i, mean + epsilon_err_high[i] + 5, f'{mean} J/m',
                ha='center', fontsize=11, fontweight='bold')

    # Annotations
    ax.text(0, 2, 'Efficient\nLeader-dominated', ha='center', fontsize=9,
            bbox=dict(boxstyle='round', facecolor='lightgreen', alpha=0.8))
    ax.text(1, 2, 'Moderate\nMixed regime', ha='center', fontsize=9,
            bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8))
    ax.text(2, 2, 'Inefficient\nStreamer-dominated', ha='center', fontsize=9,
            bbox=dict(boxstyle='round', facecolor='lightcoral', alpha=0.8))

    ax.set_ylabel('Energy Cost per Meter, epsilon (J/m)', fontsize=12)
    ax.set_xlabel('Operating Mode', fontsize=12)
    ax.set_title('Energy Cost Comparison: QCW vs Burst Mode', fontsize=14, fontweight='bold')
    ax.set_xticks(x_pos)
    ax.set_xticklabels(modes, fontsize=11)
    ax.set_yscale('log')
    ax.set_ylim(1, 150)
    ax.grid(True, alpha=0.3, axis='y')

    # Add theoretical minimum line
    ax.axhline(y=0.5, color='gray', linestyle=':', linewidth=2, alpha=0.7,
               label='Theoretical minimum (~0.5 J/m)')
    ax.legend(loc='upper left', fontsize=10)

    plt.tight_layout()
    save_figure(fig, 'epsilon-by-mode-comparison.png', 'spark-physics')


def generate_qcw_vs_burst_timeline():
    """Image 27: Side-by-side timing diagrams comparing QCW and burst operation"""
    fig, axes = plt.subplots(2, 1, figsize=(14, 8))
    fig.suptitle('QCW vs Burst Mode: Timing Comparison', fontsize=16, fontweight='bold')

    # QCW mode (top)
    ax = axes[0]
    time_qcw = np.linspace(0, 25, 1000)
    power_qcw = 50 * (time_qcw / 20)**1.5  # Gradual ramp
    power_qcw[time_qcw > 20] = 0
    length_qcw = 0.8 * (time_qcw / 20)**0.7 * 2.5  # Sub-linear growth
    length_qcw[time_qcw > 20] = length_qcw[time_qcw <= 20][-1]
    temp_qcw = 8000 + 12000 * (time_qcw / 20)  # Temperature stays high
    temp_qcw[time_qcw > 20] = 20000

    ax2 = ax.twinx()
    ax3 = ax.twinx()
    ax3.spines['right'].set_position(('outward', 60))

    p1, = ax.plot(time_qcw, power_qcw, 'r-', linewidth=2.5, label='Power (kW)')
    p2, = ax2.plot(time_qcw, length_qcw, 'b-', linewidth=2.5, label='Spark Length (m)')
    p3, = ax3.plot(time_qcw, temp_qcw, 'orange', linewidth=2.5, label='Channel Temp (K)')

    ax.set_xlabel('Time (ms)', fontsize=11)
    ax.set_ylabel('Power (kW)', fontsize=11, color='r')
    ax2.set_ylabel('Spark Length (m)', fontsize=11, color='b')
    ax3.set_ylabel('Temperature (K)', fontsize=11, color='orange')
    ax.tick_params(axis='y', labelcolor='r')
    ax2.tick_params(axis='y', labelcolor='b')
    ax3.tick_params(axis='y', labelcolor='orange')
    ax.set_title('(a) QCW Mode: 10-20 ms Ramp', fontsize=12, fontweight='bold')
    ax.grid(True, alpha=0.3)
    ax.set_xlim(0, 25)
    ax.set_ylim(0, 100)
    ax2.set_ylim(0, 3)
    ax3.set_ylim(5000, 25000)

    lines = [p1, p2, p3]
    labels = [l.get_label() for l in lines]
    ax.legend(lines, labels, loc='upper left', fontsize=9)

    ax.text(22, 70, 'Channel stays hot\nthroughout pulse', fontsize=9,
            bbox=dict(boxstyle='round', facecolor='lightgreen', alpha=0.8))

    # Burst mode (bottom)
    ax = axes[1]
    time_burst = np.linspace(0, 1, 1000)
    power_burst = np.zeros_like(time_burst)
    power_burst[(time_burst > 0.1) & (time_burst < 0.6)] = 200  # Short high pulse
    length_burst = np.zeros_like(time_burst)
    length_burst[time_burst > 0.1] = 1.2 * np.minimum((time_burst[time_burst > 0.1] - 0.1) / 0.5, 1)**0.5
    temp_burst = 8000 + 12000 * np.exp(-time_burst / 0.15)  # Rapid cooling

    ax2 = ax.twinx()
    ax3 = ax.twinx()
    ax3.spines['right'].set_position(('outward', 60))

    p1, = ax.plot(time_burst, power_burst, 'r-', linewidth=2.5, label='Power (kW)')
    p2, = ax2.plot(time_burst, length_burst, 'b-', linewidth=2.5, label='Spark Length (m)')
    p3, = ax3.plot(time_burst, temp_burst, 'orange', linewidth=2.5, label='Channel Temp (K)')

    ax.set_xlabel('Time (ms)', fontsize=11)
    ax.set_ylabel('Power (kW)', fontsize=11, color='r')
    ax2.set_ylabel('Spark Length (m)', fontsize=11, color='b')
    ax3.set_ylabel('Temperature (K)', fontsize=11, color='orange')
    ax.tick_params(axis='y', labelcolor='r')
    ax2.tick_params(axis='y', labelcolor='b')
    ax3.tick_params(axis='y', labelcolor='orange')
    ax.set_title('(b) Burst Mode: 100-500 µs Pulse', fontsize=12, fontweight='bold')
    ax.grid(True, alpha=0.3)
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 250)
    ax2.set_ylim(0, 1.5)
    ax3.set_ylim(5000, 25000)

    lines = [p1, p2, p3]
    labels = [l.get_label() for l in lines]
    ax.legend(lines, labels, loc='upper left', fontsize=9)

    ax.text(0.75, 180, 'Channel cools\nbetween pulses', fontsize=9,
            bbox=dict(boxstyle='round', facecolor='lightcoral', alpha=0.8))

    plt.tight_layout()
    save_figure(fig, 'qcw-vs-burst-timeline.png', 'spark-physics')


def generate_position_dependent_bounds():
    """Image 41: Graph showing R_min and R_max vs position"""
    fig, ax = plt.subplots(figsize=(10, 7))

    position = np.linspace(0, 1, 100)  # 0 = base, 1 = tip

    # R_min increases linearly with position
    R_min = 1e3 + (10e3 - 1e3) * position  # 1 kOhm -> 10 kOhm

    # R_max increases quadratically
    R_max = 100e3 + (100e6 - 100e3) * position**2  # 100 kOhm -> 100 MOhm

    # Typical optimized distribution (quadratic taper)
    R_opt = 5e3 + (500e3 - 5e3) * position**2

    ax.fill_between(position, R_min, R_max, alpha=0.3, color='lightblue', label='Feasible Region')
    ax.plot(position, R_min, 'b--', linewidth=2, label='R_min (Hot Leader Limit)')
    ax.plot(position, R_max, 'r--', linewidth=2, label='R_max (Cold Streamer Limit)')
    ax.plot(position, R_opt, 'g-', linewidth=3, label='Typical Optimized R', alpha=0.8)

    # Annotations
    ax.annotate('Base: Hot leader\nLow resistance', xy=(0.05, R_min[5]),
                xytext=(0.15, 2e5), fontsize=10,
                bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8),
                arrowprops=dict(arrowstyle='->', color='blue', lw=1.5))

    ax.annotate('Tip: Cold streamer\nHigh resistance', xy=(0.95, R_max[95]),
                xytext=(0.6, 5e7), fontsize=10,
                bbox=dict(boxstyle='round', facecolor='lightcoral', alpha=0.8),
                arrowprops=dict(arrowstyle='->', color='red', lw=1.5))

    ax.set_xlabel('Position Along Spark (0 = Base, 1 = Tip)', fontsize=12)
    ax.set_ylabel('Resistance (Ohm)', fontsize=12)
    ax.set_title('Position-Dependent Resistance Bounds', fontsize=14, fontweight='bold')
    ax.set_yscale('log')
    ax.set_ylim(500, 2e8)
    ax.legend(loc='upper left', fontsize=10)
    ax.grid(True, alpha=0.3, which='both')

    plt.tight_layout()
    save_figure(fig, 'position-dependent-bounds.png', 'advanced-modeling')


def generate_validation_total_resistance():
    """Image 43: Chart showing expected R_total ranges"""
    fig, ax = plt.subplots(figsize=(10, 7))

    conditions = ['Very Low Freq\n(<100 kHz)', 'Standard QCW\n(200 kHz, Leader)',
                  'Standard Burst\n(200 kHz, Streamer)', 'High Freq\n(400+ kHz)']
    R_mean = [5e3, 20e3, 150e3, 250e3]  # Mean values
    R_low = [1e3, 5e3, 50e3, 100e3]     # Lower bounds
    R_high = [10e3, 50e3, 300e3, 500e3]  # Upper bounds

    x_pos = np.arange(len(conditions))
    colors = ['green', 'lightgreen', 'orange', 'red']

    bars = ax.bar(x_pos, R_mean, color=colors, alpha=0.7, edgecolor='black', linewidth=1.5)

    # Error bars showing range
    ax.errorbar(x_pos, R_mean,
                yerr=[np.array(R_mean) - np.array(R_low), np.array(R_high) - np.array(R_mean)],
                fmt='none', color='black', capsize=10, capthick=2, linewidth=2)

    # Value labels
    for i, (mean, low, high) in enumerate(zip(R_mean, R_low, R_high)):
        ax.text(i, high * 1.3, f'{mean/1e3:.0f} kOhm\n({low/1e3:.0f}-{high/1e3:.0f})',
                ha='center', fontsize=9, fontweight='bold')

    ax.set_ylabel('Total Spark Resistance, R_total (Ohm)', fontsize=12)
    ax.set_xlabel('Operating Condition', fontsize=12)
    ax.set_title('Expected Resistance Ranges for Validation', fontsize=14, fontweight='bold')
    ax.set_xticks(x_pos)
    ax.set_xticklabels(conditions, fontsize=10)
    ax.set_yscale('log')
    ax.set_ylim(500, 1e6)
    ax.grid(True, alpha=0.3, axis='y')

    # Add note
    ax.text(0.5, 0.02, 'Note: R proportional to 1/f, R proportional to L (length)',
            transform=ax.transAxes, fontsize=10, ha='center',
            bbox=dict(boxstyle='round', facecolor='lightyellow', alpha=0.8))

    plt.tight_layout()
    save_figure(fig, 'validation-total-resistance.png', 'advanced-modeling')


def generate_lumped_vs_distributed_comparison():
    """Image 40: Table comparing lumped vs distributed models"""
    fig, ax = plt.subplots(figsize=(12, 8))
    ax.axis('off')

    # Create comparison table
    table_data = [
        ['Feature', 'Lumped Model', 'Distributed Model'],
        ['Circuit Elements', 'Single R, C_mut, C_sh', 'n segments (10-20)\nn(n+1)/2 capacitors'],
        ['Simulation Time', '~0.1 seconds', '~100-200 seconds'],
        ['Accuracy for\nShort Sparks (<1m)', 'Excellent', 'Excellent (overkill)'],
        ['Accuracy for\nLong Sparks (>2m)', 'Good approximation', 'High precision'],
        ['Spatial Detail', 'None (single point)', 'Voltage, current,\npower at each segment'],
        ['Extraction Effort', 'Simple (2x2 matrix)', 'Complex (11x11 matrix)\nPartial capacitance\ntransformation'],
        ['Best Use Cases', '- Impedance matching\n- R_opt studies\n- Quick iteration\n- Teaching', '- Research\n- Spatial distribution\n- Long sparks\n- Validation'],
        ['Computational Cost', 'Very Low', 'High'],
        ['Implementation\nDifficulty', 'Easy', 'Moderate to Hard'],
    ]

    # Color coding
    colors = []
    for i, row in enumerate(table_data):
        if i == 0:  # Header
            colors.append(['lightgray', 'lightgray', 'lightgray'])
        else:
            colors.append(['white', 'lightgreen', 'lightyellow'])

    table = ax.table(cellText=table_data, cellLoc='left', loc='center',
                     cellColours=colors, colWidths=[0.25, 0.375, 0.375])

    table.auto_set_font_size(False)
    table.set_fontsize(10)
    table.scale(1, 3)

    # Style header row
    for i in range(3):
        cell = table[(0, i)]
        cell.set_text_props(weight='bold', fontsize=11)
        cell.set_facecolor('lightblue')

    # Bold first column
    for i in range(1, len(table_data)):
        cell = table[(i, 0)]
        cell.set_text_props(weight='bold')

    ax.set_title('Lumped vs Distributed Model Comparison', fontsize=16, fontweight='bold', pad=20)

    plt.tight_layout()
    save_figure(fig, 'lumped-vs-distributed-comparison.png', 'advanced-modeling')


# ============================================================================
# MAIN EXECUTION
# ============================================================================

def generate_all_fundamentals():
    """Generate all Part 1 (Fundamentals) images"""
    print("\n" + "="*60)
    print("GENERATING PART 1: FUNDAMENTALS IMAGES")
    print("="*60)

    generate_complex_plane_admittance()  # Image 3
    generate_phase_angle_visualization()  # Image 4
    generate_phase_constraint_graph()  # Image 5
    generate_admittance_vector_addition()  # Image 7

    print(f"\n[OK] Part 1 complete: 4 images generated")


def generate_all_optimization():
    """Generate all Part 2 (Optimization) images"""
    print("\n" + "="*60)
    print("GENERATING PART 2: OPTIMIZATION IMAGES")
    print("="*60)

    generate_power_vs_resistance_curves()  # Image 9
    generate_frequency_shift_with_loading()  # Image 13
    generate_drsstc_operating_modes()  # Image 14
    generate_loaded_pole_analysis()  # Image 15

    print(f"\n[OK] Part 2 complete: 4 images generated")


def generate_all_spark_physics():
    """Generate all Part 3 (Spark Physics) images"""
    print("\n" + "="*60)
    print("GENERATING PART 3: SPARK PHYSICS IMAGES")
    print("="*60)

    generate_energy_budget_breakdown()  # Image 18
    generate_epsilon_by_mode_comparison()  # Image 19
    generate_thermal_diffusion_vs_diameter()  # Image 20
    generate_voltage_division_vs_length_plot()  # Image 24
    generate_length_vs_energy_scaling()  # Image 26
    generate_qcw_vs_burst_timeline()  # Image 27

    print(f"\n[OK] Part 3 complete: 6 images generated")


def generate_all_advanced_modeling():
    """Generate all Part 4 (Advanced Modeling) images"""
    print("\n" + "="*60)
    print("GENERATING PART 4: ADVANCED MODELING IMAGES")
    print("="*60)

    generate_capacitance_matrix_heatmap()  # Image 34
    generate_resistance_taper_initialization()  # Image 36
    generate_power_distribution_along_spark()  # Image 38
    generate_current_attenuation_plot()  # Image 39
    generate_lumped_vs_distributed_comparison()  # Image 40
    generate_position_dependent_bounds()  # Image 41
    generate_validation_total_resistance()  # Image 43

    print(f"\n[OK] Part 4 complete: 7 images generated")


def generate_all_shared():
    """Generate shared images"""
    print("\n" + "="*60)
    print("GENERATING SHARED IMAGES")
    print("="*60)

    generate_complex_number_review()  # Image 45

    print(f"\n[OK] Shared images complete: 1 image generated")


def main():
    """Main entry point"""
    parser = argparse.ArgumentParser(description='Generate Tesla Coil course images')
    parser.add_argument('--part', type=int, choices=[1, 2, 3, 4], help='Generate images for specific part only')
    parser.add_argument('--shared', action='store_true', help='Generate shared images')
    args = parser.parse_args()

    print("\n" + "="*60)
    print("TESLA COIL SPARK COURSE - IMAGE GENERATION")
    print("="*60)
    print(f"Output directories:")
    for name, path in ASSETS_DIRS.items():
        print(f"  {name}: {path}")
    print("="*60)

    if args.part:
        if args.part == 1:
            generate_all_fundamentals()
        elif args.part == 2:
            generate_all_optimization()
        elif args.part == 3:
            generate_all_spark_physics()
        elif args.part == 4:
            generate_all_advanced_modeling()
    elif args.shared:
        generate_all_shared()
    else:
        # Generate all
        generate_all_fundamentals()
        generate_all_optimization()
        generate_all_spark_physics()
        generate_all_advanced_modeling()
        generate_all_shared()

    print("\n" + "="*60)
    print("GENERATION COMPLETE!")
    print("="*60)
    print(f"\nTotal images generated: 22")
    print(f"  - Fundamentals: 4 images")
    print(f"  - Optimization: 4 images")
    print(f"  - Spark Physics: 6 images")
    print(f"  - Advanced Modeling: 7 images")
    print(f"  - Shared: 1 image")
    print("\nNote: This script generated matplotlib/numpy-based images.")
    print("Circuit diagrams (7), FEMM screenshots (5), and photos (3) require")
    print("manual creation with professional tools.")
    print("\nSee CIRCUIT-SPECIFICATIONS.md for circuit creation specs.")
    print("See IMAGE-REQUIREMENTS.md for complete list.")
    print("="*60 + "\n")


if __name__ == '__main__':
    main()