#!/usr/bin/env python3
"""RTX 5090 Power Limitation Sweep for Basalt Inference Engine.
Sweeps GPU power limits from 400W to 600W and measures prefill & decode throughput,
actual power draw, SM clocks, temperatures, and energy efficiency (Tokens/Joule).
"""

import atexit
import ctypes
import json
import os
import subprocess
import sys
import threading
import time
import urllib.request
import uuid
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

DEFAULT_POWER_W = 475
POWER_LIMITS = [400, 425, 450, 475, 500, 525, 550, 575, 600]
RUNS_PER_STEP = 2
API_URL = "http://localhost:8080/v1/chat/completions"

class NvmlSampler:
    def __init__(self, sample_interval_s=0.05):
        self.interval = sample_interval_s
        self.samples = []
        self._stop = threading.Event()
        self._thread = None
        self._ok = False
        try:
            self.nvml = ctypes.CDLL("libnvidia-ml.so.1")
            self.nvml.nvmlInit_v2()
            self.handle = ctypes.c_void_p()
            self.nvml.nvmlDeviceGetHandleByIndex_v2(0, ctypes.byref(self.handle))
            self._ok = True
        except Exception as e:
            print(f"[warn] NVML init failed: {e}", file=sys.stderr)

    def start(self):
        if not self._ok:
            return
        self.samples = []
        self._stop.clear()
        self._thread = threading.Thread(target=self._run, daemon=True)
        self._thread.start()

    def _run(self):
        pwr = ctypes.c_uint()
        clk = ctypes.c_uint()
        temp = ctypes.c_uint()
        while not self._stop.is_set():
            t = time.time()
            self.nvml.nvmlDeviceGetPowerUsage(self.handle, ctypes.byref(pwr))
            self.nvml.nvmlDeviceGetClockInfo(self.handle, 1, ctypes.byref(clk)) # SM clock
            self.nvml.nvmlDeviceGetTemperature(self.handle, 0, ctypes.byref(temp)) # GPU temp
            self.samples.append({
                "time": t,
                "power_w": pwr.value / 1000.0,
                "sm_mhz": clk.value,
                "temp_c": temp.value
            })
            time.sleep(self.interval)

    def stop(self):
        if not self._ok or self._thread is None:
            return {}
        self._stop.set()
        self._thread.join(timeout=2.0)
        if not self.samples:
            return {"avg_power_w": 0, "peak_power_w": 0, "avg_sm_mhz": 0, "max_temp_c": 0}
        powers = [s["power_w"] for s in self.samples]
        clocks = [s["sm_mhz"] for s in self.samples]
        temps = [s["temp_c"] for s in self.samples]
        return {
            "avg_power_w": round(float(np.mean(powers)), 1),
            "peak_power_w": round(float(np.max(powers)), 1),
            "avg_sm_mhz": round(float(np.mean(clocks)), 0),
            "max_temp_c": int(np.max(temps)),
            "samples_count": len(self.samples)
        }

    def close(self):
        if self._ok:
            try:
                self.nvml.nvmlShutdown()
            except Exception:
                pass


def set_gpu_power(watts):
    print(f"--> Setting GPU Power Limit to {watts}W...", flush=True)
    cmd = ["sudo", "/usr/bin/nvidia-smi", "-pl", str(watts)]
    res = subprocess.run(cmd, capture_output=True, text=True)
    if res.returncode != 0:
        print(f"Error setting power limit: {res.stderr}", file=sys.stderr)
        raise RuntimeError(res.stderr)


def post_chat(messages, max_tokens, temperature=0.0):
    body = {
        "model": "basalt",
        "messages": messages,
        "max_tokens": max_tokens,
        "temperature": temperature
    }
    req = urllib.request.Request(
        API_URL,
        headers={"Content-Type": "application/json"},
        data=json.dumps(body).encode()
    )
    t0 = time.time()
    with urllib.request.urlopen(req, timeout=120) as resp:
        data = json.loads(resp.read().decode())
    wall_s = time.time() - t0
    return data, wall_s


def run_prefill_benchmark(sampler):
    # Construct a ~16k token prompt with a unique UUID salt to bypass conversation cache
    unique_salt = f"Prompt-Run-{uuid.uuid4()}: "
    base_text = (
        "In the computational theory of large language models, matrix multiplications form the arithmetic backbone "
        "of both dense and mixture-of-experts transformer layers. As input sequences grow to tens of thousands of tokens, "
        "prefill processing shifts computational bottlenecks from memory bandwidth to tensor core arithmetic intensity. "
        "Quantized numerical representations such as NVFP4 and INT8 scale activation and weight dot-products into high-throughput tensor units. "
    )
    long_content = unique_salt + (base_text * 450)
    messages = [{"role": "user", "content": long_content}]

    sampler.start()
    data, wall_s = post_chat(messages, max_tokens=1, temperature=0.0)
    stats = sampler.stop()

    timings = data.get("timings", {})
    usage = data.get("usage", {})
    prompt_tokens = usage.get("prompt_tokens", 0)
    cached_tokens = usage.get("prompt_tokens_details", {}).get("cached_tokens", 0)
    prompt_tps = timings.get("prompt_per_second", 0.0)
    prompt_ms = timings.get("prompt_ms", wall_s * 1000.0)

    return {
        "prompt_tokens": prompt_tokens,
        "cached_tokens": cached_tokens,
        "prompt_tps": prompt_tps,
        "prompt_ms": prompt_ms,
        "wall_s": round(wall_s, 2),
        **stats
    }


def run_decode_benchmark(sampler):
    unique_salt = f"Request-{uuid.uuid4()}: "
    prompt = (
        f"{unique_salt}Write a detailed, highly technical analysis comparing speculative decoding "
        f"with standard autoregressive decoding. Cover speculative verification kernels, "
        f"draft token acceptance rates, KV cache update penalties, and hardware efficiency on modern Blackwell architecture."
    )
    messages = [{"role": "user", "content": prompt}]

    sampler.start()
    data, wall_s = post_chat(messages, max_tokens=200, temperature=0.6)
    stats = sampler.stop()

    timings = data.get("timings", {})
    usage = data.get("usage", {})
    gen_tokens = usage.get("completion_tokens", 0)
    decode_tps = timings.get("predicted_per_second", 0.0)
    decode_ms = timings.get("predicted_ms", wall_s * 1000.0)
    draft_n = timings.get("draft_n", 0)
    draft_acc = timings.get("draft_n_accepted", 0)
    acc_pct = round((draft_acc / draft_n * 100.0), 1) if draft_n > 0 else 0.0

    return {
        "gen_tokens": gen_tokens,
        "decode_tps": decode_tps,
        "decode_ms": decode_ms,
        "wall_s": round(wall_s, 2),
        "draft_n": draft_n,
        "draft_acc": draft_acc,
        "draft_acc_pct": acc_pct,
        **stats
    }


def generate_charts(results, output_png_path, output_html_path):
    # Extract arrays
    powers = [r["power_limit_w"] for r in results]
    prefill_tps = [r["prefill"]["prompt_tps"] for r in results]
    decode_tps = [r["decode"]["decode_tps"] for r in results]
    prefill_eff = [r["efficiency"]["prefill_tok_per_joule"] for r in results]
    decode_eff = [r["efficiency"]["decode_tok_per_joule"] for r in results]
    prefill_pwr = [r["prefill"]["avg_power_w"] for r in results]
    decode_pwr = [r["decode"]["avg_power_w"] for r in results]
    clocks = [r["prefill"]["avg_sm_mhz"] for r in results]
    temps = [max(r["prefill"]["max_temp_c"], r["decode"]["max_temp_c"]) for r in results]

    # Matplotlib styling
    plt.style.use("dark_background")
    fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(16, 11), dpi=200)
    fig.patch.set_facecolor("#0f172a")

    for ax in (ax1, ax2, ax3, ax4):
        ax.set_facecolor("#1e293b")
        ax.grid(True, linestyle="--", alpha=0.3, color="#94a3b8")
        ax.tick_params(colors="#cbd5e1", labelsize=10)
        for spine in ax.spines.values():
            spine.set_color("#475569")

    # 1. Prefill Throughput
    ax1.plot(powers, prefill_tps, marker="o", linewidth=2.5, markersize=7, color="#38bdf8", label="Prefill Throughput")
    for x, y in zip(powers, prefill_tps):
        ax1.annotate(f"{y:.0f}", (x, y), textcoords="offset points", xytext=(0, 8),
                     ha="center", color="#7dd3fc", fontsize=9, fontweight="bold")
    ax1.axvline(475, color="#f59e0b", linestyle=":", linewidth=1.8, label="Default Baseline (475W)")
    ax1.set_title("Prefill Throughput vs Power Limit", fontsize=13, fontweight="bold", color="#f8fafc", pad=10)
    ax1.set_xlabel("GPU Power Limit (W)", color="#cbd5e1", fontsize=11)
    ax1.set_ylabel("Tokens / Second", color="#cbd5e1", fontsize=11)
    ax1.legend(loc="lower right", facecolor="#0f172a", edgecolor="#475569")

    # 2. Decode Throughput
    ax2.plot(powers, decode_tps, marker="s", linewidth=2.5, markersize=7, color="#a855f7", label="Decode Throughput")
    for x, y in zip(powers, decode_tps):
        ax2.annotate(f"{y:.1f}", (x, y), textcoords="offset points", xytext=(0, 8),
                     ha="center", color="#c084fc", fontsize=9, fontweight="bold")
    ax2.axvline(475, color="#f59e0b", linestyle=":", linewidth=1.8, label="Default Baseline (475W)")
    ax2.set_title("Decode Throughput vs Power Limit", fontsize=13, fontweight="bold", color="#f8fafc", pad=10)
    ax2.set_xlabel("GPU Power Limit (W)", color="#cbd5e1", fontsize=11)
    ax2.set_ylabel("Tokens / Second", color="#cbd5e1", fontsize=11)
    ax2.legend(loc="lower right", facecolor="#0f172a", edgecolor="#475569")

    # 3. Energy Efficiency (Tokens per Joule)
    ax3.plot(powers, prefill_eff, marker="^", linewidth=2.5, markersize=7, color="#10b981", label="Prefill (Tok/J)")
    ax3.plot(powers, decode_eff, marker="d", linewidth=2.5, markersize=7, color="#ec4899", label="Decode (Tok/J)")
    for x, y in zip(powers, prefill_eff):
        ax3.annotate(f"{y:.2f}", (x, y), textcoords="offset points", xytext=(0, 8),
                     ha="center", color="#6ee7b7", fontsize=9, fontweight="bold")
    ax3.axvline(475, color="#f59e0b", linestyle=":", linewidth=1.8, label="Default Baseline (475W)")
    ax3.set_title("Energy Efficiency (Tokens / Joule)", fontsize=13, fontweight="bold", color="#f8fafc", pad=10)
    ax3.set_xlabel("GPU Power Limit (W)", color="#cbd5e1", fontsize=11)
    ax3.set_ylabel("Tokens per Joule", color="#cbd5e1", fontsize=11)
    ax3.legend(loc="upper right", facecolor="#0f172a", edgecolor="#475569")

    # 4. Hardware Clocks & Thermals
    ax4_t = ax4.twinx()
    p1 = ax4.plot(powers, clocks, marker="o", linewidth=2, color="#3b82f6", label="SM Clock (MHz)")
    p2 = ax4_t.plot(powers, temps, marker="x", linewidth=2, linestyle="--", color="#f43f5e", label="Peak Temp (°C)")
    ax4.set_title("Hardware Telemetry (SM Clock & Temp)", fontsize=13, fontweight="bold", color="#f8fafc", pad=10)
    ax4.set_xlabel("GPU Power Limit (W)", color="#cbd5e1", fontsize=11)
    ax4.set_ylabel("SM Clock (MHz)", color="#93c5fd", fontsize=11)
    ax4_t.set_ylabel("Temperature (°C)", color="#fda4af", fontsize=11)
    ax4_t.tick_params(colors="#fda4af", labelsize=10)
    ax4.spines["right"].set_color("#f43f5e")
    lines = p1 + p2
    labels = [l.get_label() for l in lines]
    ax4.legend(lines, labels, loc="lower right", facecolor="#0f172a", edgecolor="#475569")

    fig.suptitle(
        "NVIDIA RTX 5090 (Blackwell sm_120) — Power Sweep Benchmark (400W to 600W)\n"
        "Model: Qwen3.8-FN-NVFP4 (Basalt Local Inference Engine)",
        fontsize=15, fontweight="bold", color="#ffffff", y=0.98
    )

    plt.tight_layout(rect=[0, 0.03, 1, 0.95])
    plt.savefig(output_png_path, bbox_inches="tight", facecolor=fig.get_facecolor(), edgecolor="none")
    plt.close()
    print(f"Chart saved to {output_png_path}")

    # Generate interactive standalone HTML report
    html_content = f"""<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>RTX 5090 Power Sweep — Basalt Inference Engine</title>
<script src="https://cdn.jsdelivr.net/npm/chart.js"></script>
<style>
  :root {{
    --bg: #0b0f19;
    --card: #151d30;
    --border: #24324f;
    --text: #f1f5f9;
    --muted: #94a3b8;
    --accent: #38bdf8;
    --purple: #a855f7;
    --green: #10b981;
    --amber: #f59e0b;
  }}
  * {{ box-sizing: border-box; margin: 0; padding: 0; }}
  body {{
    font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, "Helvetica Neue", sans-serif;
    background: var(--bg);
    color: var(--text);
    padding: 2rem;
    line-height: 1.5;
  }}
  .header {{
    text-align: center;
    margin-bottom: 2rem;
  }}
  .header h1 {{
    font-size: 2rem;
    font-weight: 800;
    background: linear-gradient(135deg, #38bdf8, #a855f7);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
  }}
  .header p {{
    color: var(--muted);
    font-size: 0.95rem;
    margin-top: 0.5rem;
  }}
  .cards {{
    display: grid;
    grid-template-columns: repeat(auto-fit, minmax(220px, 1fr));
    gap: 1.25rem;
    margin-bottom: 2.5rem;
  }}
  .card {{
    background: var(--card);
    border: 1px solid var(--border);
    border-radius: 12px;
    padding: 1.25rem;
  }}
  .card .label {{
    color: var(--muted);
    font-size: 0.8rem;
    text-transform: uppercase;
    font-weight: 600;
    letter-spacing: 0.05em;
  }}
  .card .value {{
    font-size: 1.75rem;
    font-weight: 700;
    color: var(--accent);
    margin-top: 0.25rem;
  }}
  .card .sub {{
    font-size: 0.8rem;
    color: var(--muted);
    margin-top: 0.25rem;
  }}
  .charts-grid {{
    display: grid;
    grid-template-columns: repeat(auto-fit, minmax(480px, 1fr));
    gap: 1.5rem;
    margin-bottom: 2.5rem;
  }}
  .chart-container {{
    background: var(--card);
    border: 1px solid var(--border);
    border-radius: 12px;
    padding: 1.5rem;
    position: relative;
    height: 380px;
  }}
  .chart-title {{
    font-size: 1.05rem;
    font-weight: 700;
    margin-bottom: 1rem;
    color: var(--text);
  }}
  table {{
    width: 100%;
    border-collapse: collapse;
    background: var(--card);
    border-radius: 12px;
    overflow: hidden;
    border: 1px solid var(--border);
    margin-top: 1.5rem;
    font-size: 0.9rem;
  }}
  th, td {{
    padding: 0.85rem 1rem;
    text-align: right;
    border-bottom: 1px solid var(--border);
  }}
  th:first-child, td:first-child {{
    text-align: left;
  }}
  th {{
    background: #1b2640;
    color: #cbd5e1;
    font-weight: 600;
  }}
  tr:hover {{
    background: #1e2c4a;
  }}
  .highlight {{
    color: var(--amber);
    font-weight: 700;
  }}
</style>
</head>
<body>
  <div class="header">
    <h1>NVIDIA GeForce RTX 5090 Power Scaling Sweep</h1>
    <p>Engine: <strong>Basalt (sm_120 Blackwell)</strong> &bull; Model: <strong>Qwen3.8-FN-NVFP4 (122GB)</strong> &bull; Sweep Range: 400W &ndash; 600W</p>
  </div>

  <div class="cards">
    <div class="card">
      <div class="label">Optimal Efficiency Point</div>
      <div class="value" id="optEff">-</div>
      <div class="sub">Best Tokens / Joule balance</div>
    </div>
    <div class="card">
      <div class="label">Prefill Range</div>
      <div class="value" id="prefillRange">-</div>
      <div class="sub">400W vs 600W scaling</div>
    </div>
    <div class="card">
      <div class="label">Decode Range</div>
      <div class="value" id="decodeRange">-</div>
      <div class="sub">400W vs 600W scaling</div>
    </div>
    <div class="card">
      <div class="label">Baseline (475W)</div>
      <div class="value" id="baselinePerf">-</div>
      <div class="sub">User's current profile</div>
    </div>
  </div>

  <div class="charts-grid">
    <div class="chart-container">
      <div class="chart-title">Prefill &amp; Decode Throughput (Tokens/s)</div>
      <canvas id="throughputChart"></canvas>
    </div>
    <div class="chart-container">
      <div class="chart-title">Energy Efficiency (Tokens / Joule)</div>
      <canvas id="efficiencyChart"></canvas>
    </div>
    <div class="chart-container">
      <div class="chart-title">Measured Power Draw vs Power Limit (W)</div>
      <canvas id="powerChart"></canvas>
    </div>
    <div class="chart-container">
      <div class="chart-title">SM Clocks (MHz) &amp; Peak Thermals (°C)</div>
      <canvas id="hardwareChart"></canvas>
    </div>
  </div>

  <h2 style="font-size: 1.25rem; font-weight: 700; margin-top: 2rem;">Full Sweep Results</h2>
  <table>
    <thead>
      <tr>
        <th>Power Limit</th>
        <th>Actual Prefill W</th>
        <th>Prefill TPS</th>
        <th>Prefill Tok/J</th>
        <th>Actual Decode W</th>
        <th>Decode TPS</th>
        <th>Decode Tok/J</th>
        <th>SM Clock</th>
        <th>Peak Temp</th>
      </tr>
    </thead>
    <tbody>
"""
    base_row = next((r for r in results if r["power_limit_w"] == 475), results[len(results)//2])
    for r in results:
        is_base = (r["power_limit_w"] == 475)
        cls = ' class="highlight"' if is_base else ""
        star = " ★ (Baseline)" if is_base else ""
        html_content += f"""      <tr{cls}>
        <td><strong>{r["power_limit_w"]}W</strong>{star}</td>
        <td>{r["prefill"]["avg_power_w"]} W</td>
        <td><strong>{r["prefill"]["prompt_tps"]:.1f}</strong> t/s</td>
        <td>{r["efficiency"]["prefill_tok_per_joule"]:.2f}</td>
        <td>{r["decode"]["avg_power_w"]} W</td>
        <td><strong>{r["decode"]["decode_tps"]:.1f}</strong> t/s</td>
        <td>{r["efficiency"]["decode_tok_per_joule"]:.2f}</td>
        <td>{r["prefill"]["avg_sm_mhz"]:.0f} MHz</td>
        <td>{max(r["prefill"]["max_temp_c"], r["decode"]["max_temp_c"])} °C</td>
      </tr>
"""

    html_content += f"""    </tbody>
  </table>

<script>
  const data = {json.dumps(results)};
  const labels = data.map(d => d.power_limit_w + "W");
  const prefillTps = data.map(d => d.prefill.prompt_tps);
  const decodeTps = data.map(d => d.decode.decode_tps);
  const prefillEff = data.map(d => d.efficiency.prefill_tok_per_joule);
  const decodeEff = data.map(d => d.efficiency.decode_tok_per_joule);
  const prefillPwr = data.map(d => d.prefill.avg_power_w);
  const decodePwr = data.map(d => d.decode.avg_power_w);
  const smClocks = data.map(d => d.prefill.avg_sm_mhz);
  const temps = data.map(d => Math.max(d.prefill.max_temp_c, d.decode.max_temp_c));

  // Summary cards
  const b = data.find(d => d.power_limit_w === 475) || data[0];
  document.getElementById("baselinePerf").innerText = b.decode.decode_tps.toFixed(1) + " t/s";
  document.getElementById("prefillRange").innerText = data[0].prefill.prompt_tps.toFixed(0) + " → " + data[data.length-1].prefill.prompt_tps.toFixed(0);
  document.getElementById("decodeRange").innerText = data[0].decode.decode_tps.toFixed(1) + " → " + data[data.length-1].decode.decode_tps.toFixed(1);

  // Best efficiency
  let maxEffIdx = 0;
  for(let i=1; i<data.length; ++i) {{
    if(data[i].efficiency.prefill_tok_per_joule > data[maxEffIdx].efficiency.prefill_tok_per_joule) maxEffIdx = i;
  }}
  document.getElementById("optEff").innerText = data[maxEffIdx].power_limit_w + "W (" + data[maxEffIdx].efficiency.prefill_tok_per_joule.toFixed(2) + " T/J)";

  Chart.defaults.color = "#94a3b8";
  Chart.defaults.borderColor = "#24324f";

  new Chart(document.getElementById("throughputChart"), {{
    type: "line",
    data: {{
      labels: labels,
      datasets: [
        {{ label: "Prefill TPS", data: prefillTps, borderColor: "#38bdf8", backgroundColor: "#38bdf820", tension: 0.3, yAxisID: "y" }},
        {{ label: "Decode TPS", data: decodeTps, borderColor: "#a855f7", backgroundColor: "#a855f720", tension: 0.3, yAxisID: "y1" }}
      ]
    }},
    options: {{
      responsive: true, maintainAspectRatio: false,
      scales: {{
        y: {{ type: "linear", display: true, position: "left", title: {{ display: true, text: "Prefill TPS", color: "#38bdf8" }} }},
        y1: {{ type: "linear", display: true, position: "right", title: {{ display: true, text: "Decode TPS", color: "#a855f7" }}, grid: {{ drawOnChartArea: false }} }}
      }}
    }}
  }});

  new Chart(document.getElementById("efficiencyChart"), {{
    type: "line",
    data: {{
      labels: labels,
      datasets: [
        {{ label: "Prefill (Tok/Joule)", data: prefillEff, borderColor: "#10b981", tension: 0.3 }},
        {{ label: "Decode (Tok/Joule)", data: decodeEff, borderColor: "#ec4899", tension: 0.3 }}
      ]
    }},
    options: {{
      responsive: true, maintainAspectRatio: false,
      scales: {{ y: {{ title: {{ display: true, text: "Tokens per Joule" }} }} }}
    }}
  }});

  new Chart(document.getElementById("powerChart"), {{
    type: "bar",
    data: {{
      labels: labels,
      datasets: [
        {{ label: "Prefill Avg Draw (W)", data: prefillPwr, backgroundColor: "#38bdf8aa" }},
        {{ label: "Decode Avg Draw (W)", data: decodePwr, backgroundColor: "#a855f7aa" }}
      ]
    }},
    options: {{
      responsive: true, maintainAspectRatio: false,
      scales: {{ y: {{ title: {{ display: true, text: "Power Draw (Watts)" }} }} }}
    }}
  }});

  new Chart(document.getElementById("hardwareChart"), {{
    type: "line",
    data: {{
      labels: labels,
      datasets: [
        {{ label: "SM Clock (MHz)", data: smClocks, borderColor: "#3b82f6", tension: 0.3, yAxisID: "y" }},
        {{ label: "Peak Temp (°C)", data: temps, borderColor: "#f43f5e", tension: 0.3, yAxisID: "y1" }}
      ]
    }},
    options: {{
      responsive: true, maintainAspectRatio: false,
      scales: {{
        y: {{ title: {{ display: true, text: "Clock (MHz)", color: "#3b82f6" }} }},
        y1: {{ position: "right", title: {{ display: true, text: "Temperature (°C)", color: "#f43f5e" }}, grid: {{ drawOnChartArea: false }} }}
      }}
    }}
  }});
</script>
</body>
</html>
"""
    with open(output_html_path, "w", encoding="utf-8") as f:
        f.write(html_content)
    print(f"Interactive HTML report saved to {output_html_path}")


def main():
    sampler = NvmlSampler()
    atexit.register(lambda: (set_gpu_power(DEFAULT_POWER_W), sampler.close()))

    print(f"=== Starting RTX 5090 Power Sweep ({POWER_LIMITS[0]}W - {POWER_LIMITS[-1]}W) ===")
    results = []

    try:
        for pwr in POWER_LIMITS:
            set_gpu_power(pwr)
            time.sleep(1.5) # Let power and clocks settle

            # Run 2 passes to average out any transient jitter
            prefill_runs = []
            decode_runs = []

            for r_idx in range(RUNS_PER_STEP):
                print(f"  [{pwr}W] Run {r_idx+1}/{RUNS_PER_STEP}: Prefill test (~16k tokens)...", end="", flush=True)
                p_res = run_prefill_benchmark(sampler)
                print(f" -> {p_res['prompt_tps']:.1f} t/s ({p_res['avg_power_w']}W)", flush=True)
                prefill_runs.append(p_res)
                time.sleep(0.5)

                print(f"  [{pwr}W] Run {r_idx+1}/{RUNS_PER_STEP}: Decode test (~200 tokens)...", end="", flush=True)
                d_res = run_decode_benchmark(sampler)
                print(f" -> {d_res['decode_tps']:.1f} t/s ({d_res['avg_power_w']}W)", flush=True)
                decode_runs.append(d_res)
                time.sleep(0.5)

            # Average metrics across runs
            avg_p_tps = float(np.mean([x["prompt_tps"] for x in prefill_runs]))
            avg_p_pwr = float(np.mean([x["avg_power_w"] for x in prefill_runs]))
            avg_p_sm = float(np.mean([x["avg_sm_mhz"] for x in prefill_runs]))
            max_p_temp = int(np.max([x["max_temp_c"] for x in prefill_runs]))

            avg_d_tps = float(np.mean([x["decode_tps"] for x in decode_runs]))
            avg_d_pwr = float(np.mean([x["avg_power_w"] for x in decode_runs]))
            avg_d_sm = float(np.mean([x["avg_sm_mhz"] for x in decode_runs]))
            max_d_temp = int(np.max([x["max_temp_c"] for x in decode_runs]))
            avg_draft_acc = float(np.mean([x["draft_acc_pct"] for x in decode_runs]))

            eff_p = avg_p_tps / avg_p_pwr if avg_p_pwr > 0 else 0.0
            eff_d = avg_d_tps / avg_d_pwr if avg_d_pwr > 0 else 0.0

            step_data = {
                "power_limit_w": pwr,
                "prefill": {
                    "prompt_tps": round(avg_p_tps, 1),
                    "avg_power_w": round(avg_p_pwr, 1),
                    "avg_sm_mhz": round(avg_p_sm, 0),
                    "max_temp_c": max_p_temp,
                    "runs": prefill_runs
                },
                "decode": {
                    "decode_tps": round(avg_d_tps, 1),
                    "avg_power_w": round(avg_d_pwr, 1),
                    "avg_sm_mhz": round(avg_d_sm, 0),
                    "max_temp_c": max_d_temp,
                    "draft_acc_pct": round(avg_draft_acc, 1),
                    "runs": decode_runs
                },
                "efficiency": {
                    "prefill_tok_per_joule": round(eff_p, 3),
                    "decode_tok_per_joule": round(eff_d, 3),
                    "prefill_joules_per_tok": round(avg_p_pwr / avg_p_tps, 4) if avg_p_tps else 0,
                    "decode_joules_per_tok": round(avg_d_pwr / avg_d_tps, 4) if avg_d_tps else 0
                }
            }
            results.append(step_data)
            print(f"==> [{pwr}W Summary] Prefill: {avg_p_tps:.1f} t/s ({eff_p:.2f} tok/J) | Decode: {avg_d_tps:.1f} t/s ({eff_d:.2f} tok/J)\n")

    finally:
        set_gpu_power(DEFAULT_POWER_W)
        sampler.close()

    # Save JSON results
    json_path = "scripts/power_sweep_results.json"
    with open(json_path, "w", encoding="utf-8") as f:
        json.dump(results, f, indent=2)
    print(f"Results written to {json_path}")

    # Generate visual outputs
    png_path = "scripts/power_sweep_chart.png"
    html_path = "scripts/power_sweep_report.html"
    generate_charts(results, png_path, html_path)
    print("=== Sweep Complete ===")


if __name__ == "__main__":
    main()