#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
GPU Tune Tool
=============

支持功能：
1. benchmark：遍历功率 + 风扇组合，寻找温度/性能平衡点
2. auto：根据温度自动调节风扇和功率
3. set：手动设置功率和风扇
4. restore：恢复自动风扇控制
5. status：查看当前状态

适配环境：
- 8 张 NVIDIA GeForce RTX 5090
- 每张卡 3 个风扇
- 显卡输入编号为 1~8
- Xorg DISPLAY 默认为 :0

示例：
    python3 gpu_tune_tool.py status --gpus all

    python3 gpu_tune_tool.py set --gpus 1,2,5 --power 500 --fan 80

    python3 gpu_tune_tool.py restore --gpus all

    python3 gpu_tune_tool.py auto \
        --gpus all \
        --target-power 500 \
        --min-power 350 \
        --max-fan 95 \
        --min-fan 60 \
        --hot-temp 78 \
        --cool-temp 68

    python3 gpu_tune_tool.py benchmark \
        --gpus all \
        --time 10h \
        --stress "./gpu_burn 36000" \
        --power-limits "350,400,450,500,550" \
        --fan-speeds "60,70,80,90"
"""

import os
import re
import csv
import sys
import time
import signal
import argparse
import subprocess
import statistics
from pathlib import Path
from datetime import datetime, timedelta
from typing import List, Optional, Dict, Any


TOTAL_GPUS = 8
FANS_PER_GPU = 3

DEFAULT_POWER_LIMITS = [350, 400, 450, 500, 550]
DEFAULT_FAN_SPEEDS = [60, 70, 80, 90]

DEFAULT_WARMUP_SECONDS = 120
DEFAULT_MONITOR_INTERVAL = 10


# ============================================================
# Common Utils
# ============================================================

def ensure_display():
    if not os.environ.get("DISPLAY"):
        os.environ["DISPLAY"] = ":0"


def run_cmd(cmd: str, check: bool = True) -> str:
    r = subprocess.run(cmd, shell=True, capture_output=True, text=True)
    if check and r.returncode != 0:
        print(f"[ERROR] command failed: {cmd}")
        if r.stdout.strip():
            print(f"stdout: {r.stdout.strip()}")
        if r.stderr.strip():
            print(f"stderr: {r.stderr.strip()}")
        raise RuntimeError(cmd)
    return r.stdout.strip()


def parse_time(s: str) -> int:
    m = re.match(r"^(\d+)(h|m|s)?$", s)
    if not m:
        raise ValueError(f"Invalid time format: {s}")
    value = int(m.group(1))
    unit = m.group(2) or "s"
    return value * {"h": 3600, "m": 60, "s": 1}[unit]


def parse_int_list(s: str) -> List[int]:
    return [int(x.strip()) for x in s.split(",") if x.strip()]


def expand_gpu_list(spec: str) -> List[int]:
    if spec.strip().lower() == "all":
        return list(range(1, TOTAL_GPUS + 1))

    result = []
    for part in spec.split(","):
        part = part.strip()
        if not part:
            continue
        if not part.isdigit():
            raise ValueError(f"Invalid GPU number: {part}")
        n = int(part)
        if n < 1 or n > TOTAL_GPUS:
            raise ValueError(f"GPU {n} out of range 1-{TOTAL_GPUS}")
        result.append(n)

    return sorted(set(result))


def user_gpu_to_id(user_gpu: int) -> int:
    return user_gpu - 1


def gpu_fan_ids(user_gpu: int) -> List[int]:
    gpu_id = user_gpu_to_id(user_gpu)
    start = gpu_id * FANS_PER_GPU
    return list(range(start, start + FANS_PER_GPU))


# ============================================================
# NVIDIA Control
# ============================================================

def set_power_limit(user_gpu: int, watts: int):
    gpu_id = user_gpu_to_id(user_gpu)
    print(f"[SET] GPU-{user_gpu} power limit -> {watts}W")
    run_cmd(f"nvidia-smi -i {gpu_id} -pl {watts}")


def enable_manual_fan(user_gpu: int):
    gpu_id = user_gpu_to_id(user_gpu)
    run_cmd(f'nvidia-settings -a "[gpu:{gpu_id}]/GPUFanControlState=1"')


def disable_manual_fan(user_gpu: int):
    gpu_id = user_gpu_to_id(user_gpu)
    run_cmd(f'nvidia-settings -a "[gpu:{gpu_id}]/GPUFanControlState=0"', check=False)


def set_fan_speed(user_gpu: int, speed: int):
    gpu_id = user_gpu_to_id(user_gpu)
    fans = gpu_fan_ids(user_gpu)

    print(f"[SET] GPU-{user_gpu} fan -> {speed}%  fans={fans}")
    enable_manual_fan(user_gpu)

    for fan_id in fans:
        run_cmd(f'nvidia-settings -a "[fan:{fan_id}]/GPUTargetFanSpeed={speed}"')


def query_gpu(user_gpu: int) -> Dict[str, Any]:
    gpu_id = user_gpu_to_id(user_gpu)

    fields = (
        "timestamp,index,name,temperature.gpu,power.draw,power.limit,"
        "utilization.gpu,utilization.memory,memory.used,memory.total,"
        "clocks.sm,clocks.mem,pstate"
    )

    out = run_cmd(
        f"nvidia-smi --query-gpu={fields} "
        f"--format=csv,noheader,nounits -i {gpu_id}"
    )

    parts = [p.strip() for p in out.split(",")]

    return {
        "timestamp": parts[0],
        "user_gpu": user_gpu,
        "gpu_index": int(parts[1]),
        "name": parts[2],
        "temperature": float(parts[3]),
        "power_draw": float(parts[4]),
        "power_limit": float(parts[5]),
        "gpu_util": float(parts[6]),
        "mem_util": float(parts[7]),
        "memory_used": float(parts[8]),
        "memory_total": float(parts[9]),
        "clocks_sm": float(parts[10]),
        "clocks_mem": float(parts[11]),
        "pstate": parts[12],
    }


def query_fan_speed_percent(user_gpu: int) -> Optional[float]:
    """
    查询该 GPU 第一个 fan 的目标转速。
    对你的机器来说每张卡 3 个 fan，脚本设置时三个 fan 会保持一致。
    """
    fan_id = gpu_fan_ids(user_gpu)[0]
    out = run_cmd(
        f'nvidia-settings -q "[fan:{fan_id}]/GPUTargetFanSpeed"',
        check=False
    )

    m = re.search(r":\s*([0-9]+)\.", out)
    if m:
        return float(m.group(1))

    m = re.search(r":\s*([0-9]+)", out)
    if m:
        return float(m.group(1))

    return None


def query_power_limits(user_gpu: int) -> Dict[str, Optional[float]]:
    """
    尝试查询 Min/Max/Default/Current Power Limit。
    不同驱动输出格式略有差异，所以这里做宽松解析。
    """
    gpu_id = user_gpu_to_id(user_gpu)
    out = run_cmd(f"nvidia-smi -i {gpu_id} -q -d POWER", check=False)

    result = {
        "min": None,
        "max": None,
        "default": None,
        "current": None,
    }

    patterns = {
        "min": r"Min Power Limit\s*:\s*([0-9.]+)\s*W",
        "max": r"Max Power Limit\s*:\s*([0-9.]+)\s*W",
        "default": r"Default Power Limit\s*:\s*([0-9.]+)\s*W",
        "current": r"Current Power Limit\s*:\s*([0-9.]+)\s*W",
    }

    for key, pat in patterns.items():
        m = re.search(pat, out)
        if m:
            result[key] = float(m.group(1))

    return result


# ============================================================
# Modes
# ============================================================

def mode_status(args):
    gpus = expand_gpu_list(args.gpus)

    print()
    print("GPU Status")
    print("-" * 110)
    print(
        f"{'GPU':<5} {'Temp':<7} {'Power':<12} {'Limit':<8} "
        f"{'FanTarget':<10} {'Util':<8} {'Mem':<18} {'SMClock':<10} {'PState':<8}"
    )
    print("-" * 110)

    for g in gpus:
        try:
            d = query_gpu(g)
            fan = query_fan_speed_percent(g)
            fan_str = f"{fan:.0f}%" if fan is not None else "N/A"

            print(
                f"{g:<5} "
                f"{d['temperature']:<7.1f} "
                f"{d['power_draw']:<6.1f}W     "
                f"{d['power_limit']:<8.0f} "
                f"{fan_str:<10} "
                f"{d['gpu_util']:<8.1f} "
                f"{d['memory_used']:.0f}/{d['memory_total']:.0f} MB   "
                f"{d['clocks_sm']:<10.0f} "
                f"{d['pstate']:<8}"
            )
        except Exception as e:
            print(f"GPU-{g}: query failed: {e}")


def mode_set(args):
    gpus = expand_gpu_list(args.gpus)

    if args.power is None and args.fan is None:
        raise ValueError("set mode requires --power or --fan")

    for g in gpus:
        if args.power is not None:
            set_power_limit(g, args.power)

        if args.fan is not None:
            set_fan_speed(g, args.fan)

    print("[DONE] set completed.")


def mode_restore(args):
    gpus = expand_gpu_list(args.gpus)

    for g in gpus:
        print(f"[RESTORE] GPU-{g} fan auto")
        disable_manual_fan(g)

    print("[DONE] restore completed.")


# ============================================================
# Auto Mode
# ============================================================

def auto_control_once(
    user_gpu: int,
    state: Dict[int, Dict[str, Any]],
    target_power: int,
    min_power: int,
    max_power: int,
    min_fan: int,
    max_fan: int,
    hot_temp: float,
    cool_temp: float,
    critical_temp: float,
    fan_step: int,
    power_step: int,
):
    d = query_gpu(user_gpu)

    if user_gpu not in state:
        state[user_gpu] = {
            "current_power": int(d["power_limit"]),
            "current_fan": query_fan_speed_percent(user_gpu) or min_fan,
        }

    current_power = int(state[user_gpu]["current_power"])
    current_fan = int(state[user_gpu]["current_fan"])

    temp = d["temperature"]

    action = "hold"

    # 极端高温：风扇直接拉满，并快速降功率
    if temp >= critical_temp:
        new_fan = max_fan
        new_power = max(min_power, current_power - power_step * 2)

        if new_fan != current_fan:
            set_fan_speed(user_gpu, new_fan)
            current_fan = new_fan

        if new_power != current_power:
            set_power_limit(user_gpu, new_power)
            current_power = new_power

        action = f"CRITICAL: fan->{current_fan}%, power->{current_power}W"

    # 高温：先加风扇，风扇没空间再降功率
    elif temp >= hot_temp:
        if current_fan < max_fan:
            new_fan = min(max_fan, current_fan + fan_step)
            set_fan_speed(user_gpu, new_fan)
            current_fan = new_fan
            action = f"hot: fan->{current_fan}%"
        else:
            new_power = max(min_power, current_power - power_step)
            if new_power < current_power:
                set_power_limit(user_gpu, new_power)
                current_power = new_power
                action = f"hot: power->{current_power}W"
            else:
                action = "hot: already min power"

    # 低温：先恢复功率，再慢慢降低风扇
    elif temp <= cool_temp:
        if current_power < target_power:
            new_power = min(target_power, current_power + power_step)
            set_power_limit(user_gpu, new_power)
            current_power = new_power
            action = f"cool: power->{current_power}W"
        elif current_fan > min_fan:
            new_fan = max(min_fan, current_fan - fan_step)
            set_fan_speed(user_gpu, new_fan)
            current_fan = new_fan
            action = f"cool: fan->{current_fan}%"
        else:
            action = "cool: already target"

    state[user_gpu]["current_power"] = current_power
    state[user_gpu]["current_fan"] = current_fan

    return {
        "timestamp": d["timestamp"],
        "gpu": user_gpu,
        "temperature": temp,
        "power_draw": d["power_draw"],
        "power_limit": current_power,
        "fan": current_fan,
        "gpu_util": d["gpu_util"],
        "clocks_sm": d["clocks_sm"],
        "action": action,
    }


def mode_auto(args):
    gpus = expand_gpu_list(args.gpus)

    out_dir = Path(f"gpu-auto-result-{datetime.now().strftime('%Y%m%d-%H%M%S')}")
    out_dir.mkdir(parents=True, exist_ok=True)

    log_path = out_dir / "auto_log.csv"

    stop_flag = False

    def handle_signal(sig, frame):
        nonlocal stop_flag
        stop_flag = True

    signal.signal(signal.SIGINT, handle_signal)
    signal.signal(signal.SIGTERM, handle_signal)

    duration = parse_time(args.time) if args.time else None
    start_time = time.time()

    state: Dict[int, Dict[str, Any]] = {}

    print()
    print("Auto Control Mode")
    print("-" * 80)
    print(f"GPUs:          {gpus}")
    print(f"Duration:      {args.time or 'unlimited'}")
    print(f"Target power:  {args.target_power}W")
    print(f"Min power:     {args.min_power}W")
    print(f"Max power:     {args.max_power}W")
    print(f"Fan range:     {args.min_fan}% - {args.max_fan}%")
    print(f"Hot temp:      {args.hot_temp}C")
    print(f"Cool temp:     {args.cool_temp}C")
    print(f"Critical temp: {args.critical_temp}C")
    print(f"Interval:      {args.interval}s")
    print(f"Log:           {log_path}")
    print("-" * 80)

    # 初始设置：目标功率 + 最低风扇
    for g in gpus:
        set_power_limit(g, min(args.target_power, args.max_power))
        set_fan_speed(g, args.min_fan)
        state[g] = {
            "current_power": min(args.target_power, args.max_power),
            "current_fan": args.min_fan,
        }

    with open(log_path, "w", newline="") as f:
        fields = [
            "timestamp",
            "gpu",
            "temperature",
            "power_draw",
            "power_limit",
            "fan",
            "gpu_util",
            "clocks_sm",
            "action",
        ]
        writer = csv.DictWriter(f, fieldnames=fields)
        writer.writeheader()

        try:
            while not stop_flag:
                if duration is not None and time.time() - start_time >= duration:
                    break

                for g in gpus:
                    try:
                        row = auto_control_once(
                            user_gpu=g,
                            state=state,
                            target_power=min(args.target_power, args.max_power),
                            min_power=args.min_power,
                            max_power=args.max_power,
                            min_fan=args.min_fan,
                            max_fan=args.max_fan,
                            hot_temp=args.hot_temp,
                            cool_temp=args.cool_temp,
                            critical_temp=args.critical_temp,
                            fan_step=args.fan_step,
                            power_step=args.power_step,
                        )

                        writer.writerow(row)
                        f.flush()

                        print(
                            f"[GPU-{row['gpu']}] "
                            f"T={row['temperature']:.1f}C "
                            f"Pdraw={row['power_draw']:.1f}W "
                            f"Plimit={row['power_limit']}W "
                            f"Fan={row['fan']}% "
                            f"Util={row['gpu_util']:.0f}% "
                            f"Clock={row['clocks_sm']:.0f}MHz "
                            f"=> {row['action']}"
                        )

                    except Exception as e:
                        print(f"[WARN] GPU-{g} auto control failed: {e}")

                time.sleep(args.interval)

        finally:
            if args.restore_on_exit:
                print("[EXIT] restore fan auto...")
                for g in gpus:
                    disable_manual_fan(g)

            print(f"[DONE] auto log saved: {log_path}")


# ============================================================
# Benchmark Mode
# ============================================================

class StressProcess:
    def __init__(self, cmd: str):
        self.cmd = cmd
        self.proc: Optional[subprocess.Popen] = None

    def start(self):
        print(f"[STRESS START] {self.cmd}")
        self.proc = subprocess.Popen(
            self.cmd,
            shell=True,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
            preexec_fn=os.setsid,
        )

    def stop(self):
        if self.proc and self.proc.poll() is None:
            print(f"[STRESS STOP] pid={self.proc.pid}")
            try:
                os.killpg(os.getpgid(self.proc.pid), signal.SIGTERM)
                self.proc.wait(timeout=10)
            except Exception:
                try:
                    os.killpg(os.getpgid(self.proc.pid), signal.SIGKILL)
                except Exception:
                    pass


def mode_benchmark(args):
    gpus = expand_gpu_list(args.gpus)
    total_seconds = parse_time(args.time)
    power_limits = parse_int_list(args.power_limits)
    fan_speeds = parse_int_list(args.fan_speeds)

    combo_count = len(power_limits) * len(fan_speeds)
    secs_per_combo = total_seconds // combo_count

    if secs_per_combo <= args.warmup + 60:
        raise ValueError(
            f"Total time too short. combo={combo_count}, "
            f"each={secs_per_combo}s, warmup={args.warmup}s"
        )

    sample_seconds = secs_per_combo - args.warmup

    out_dir = Path(f"gpu-benchmark-result-{datetime.now().strftime('%Y%m%d-%H%M%S')}")
    out_dir.mkdir(parents=True, exist_ok=True)

    csv_path = out_dir / "result.csv"

    print()
    print("Benchmark Mode")
    print("-" * 80)
    print(f"GPUs:           {gpus}")
    print(f"Total time:     {args.time}")
    print(f"Power limits:   {power_limits}")
    print(f"Fan speeds:     {fan_speeds}")
    print(f"Combinations:   {combo_count}")
    print(f"Per combo:      {secs_per_combo}s")
    print(f"Warmup:         {args.warmup}s")
    print(f"Sampling:       {sample_seconds}s")
    print(f"Stress command: {args.stress}")
    print(f"Output:         {csv_path}")
    print("-" * 80)

    stress = StressProcess(args.stress)
    stress.start()

    stop_flag = False

    def handle_signal(sig, frame):
        nonlocal stop_flag
        stop_flag = True

    signal.signal(signal.SIGINT, handle_signal)
    signal.signal(signal.SIGTERM, handle_signal)

    fields = [
        "group_id",
        "power_limit_w",
        "fan_speed_pct",
        "timestamp",
        "gpu",
        "gpu_index",
        "temperature_c",
        "power_draw_w",
        "gpu_util_pct",
        "mem_util_pct",
        "memory_used_mb",
        "memory_total_mb",
        "clocks_sm_mhz",
        "clocks_mem_mhz",
        "pstate",
    ]

    try:
        with open(csv_path, "w", newline="") as f:
            writer = csv.DictWriter(f, fieldnames=fields)
            writer.writeheader()

            group_id = 0

            for power in power_limits:
                for fan in fan_speeds:
                    if stop_flag:
                        break

                    group_id += 1
                    print()
                    print(f"[GROUP {group_id}/{combo_count}] power={power}W fan={fan}%")

                    for g in gpus:
                        set_power_limit(g, power)
                        set_fan_speed(g, fan)

                    print(f"[WARMUP] {args.warmup}s")
                    time.sleep(args.warmup)

                    print(f"[SAMPLING] {sample_seconds}s")
                    t0 = time.time()

                    while time.time() - t0 < sample_seconds and not stop_flag:
                        for g in gpus:
                            d = query_gpu(g)

                            writer.writerow({
                                "group_id": group_id,
                                "power_limit_w": power,
                                "fan_speed_pct": fan,
                                "timestamp": d["timestamp"],
                                "gpu": g,
                                "gpu_index": d["gpu_index"],
                                "temperature_c": d["temperature"],
                                "power_draw_w": d["power_draw"],
                                "gpu_util_pct": d["gpu_util"],
                                "mem_util_pct": d["mem_util"],
                                "memory_used_mb": d["memory_used"],
                                "memory_total_mb": d["memory_total"],
                                "clocks_sm_mhz": d["clocks_sm"],
                                "clocks_mem_mhz": d["clocks_mem"],
                                "pstate": d["pstate"],
                            })

                        f.flush()
                        time.sleep(args.interval)

    finally:
        stress.stop()
        if args.restore_on_exit:
            for g in gpus:
                disable_manual_fan(g)

    analyze_benchmark(csv_path)


def analyze_benchmark(csv_path: Path):
    groups: Dict[int, Dict[str, Any]] = {}

    with open(csv_path, "r") as f:
        reader = csv.DictReader(f)

        for row in reader:
            gid = int(row["group_id"])

            if gid not in groups:
                groups[gid] = {
                    "power": int(row["power_limit_w"]),
                    "fan": int(row["fan_speed_pct"]),
                    "temps": [],
                    "powers": [],
                    "utils": [],
                    "clocks": [],
                }

            groups[gid]["temps"].append(float(row["temperature_c"]))
            groups[gid]["powers"].append(float(row["power_draw_w"]))
            groups[gid]["utils"].append(float(row["gpu_util_pct"]))
            groups[gid]["clocks"].append(float(row["clocks_sm_mhz"]))

    if not groups:
        print("[WARN] no benchmark data.")
        return

    summaries = []

    for gid, d in groups.items():
        summaries.append({
            "group": gid,
            "power": d["power"],
            "fan": d["fan"],
            "avg_temp": statistics.mean(d["temps"]),
            "max_temp": max(d["temps"]),
            "avg_power": statistics.mean(d["powers"]),
            "avg_util": statistics.mean(d["utils"]),
            "avg_clock": statistics.mean(d["clocks"]),
        })

    max_clock = max(x["avg_clock"] for x in summaries) or 1
    max_temp = max(x["avg_temp"] for x in summaries) or 1

    for s in summaries:
        perf_score = s["avg_clock"] / max_clock
        temp_penalty = s["avg_temp"] / max_temp
        power_penalty = s["avg_power"] / s["power"]
        s["score"] = perf_score - 0.3 * temp_penalty - 0.2 * power_penalty

    summaries.sort(key=lambda x: x["score"], reverse=True)

    print()
    print("Benchmark Ranking")
    print("-" * 90)
    print(
        f"{'Rank':<6} {'Power':<8} {'Fan':<6} {'AvgTemp':<9} "
        f"{'MaxTemp':<9} {'AvgPower':<10} {'AvgUtil':<9} {'AvgClock':<10} {'Score':<8}"
    )
    print("-" * 90)

    for i, s in enumerate(summaries, 1):
        mark = "*" if i <= 3 else ""
        print(
            f"{i:<6} {s['power']:<8} {s['fan']:<6} "
            f"{s['avg_temp']:<9.1f} {s['max_temp']:<9.1f} "
            f"{s['avg_power']:<10.1f} {s['avg_util']:<9.1f} "
            f"{s['avg_clock']:<10.0f} {s['score']:<8.3f} {mark}"
        )

    best = summaries[0]
    print()
    print(f"Recommended: Power={best['power']}W, Fan={best['fan']}%")


# ============================================================
# Argparse
# ============================================================

def build_parser():
    parser = argparse.ArgumentParser(
        description="GPU power/fan tuning tool"
    )

    sub = parser.add_subparsers(dest="mode", required=True)

    # status
    p = sub.add_parser("status", help="show GPU status")
    p.add_argument("--gpus", required=True, help="all or 1,2,5")
    p.set_defaults(func=mode_status)

    # set
    p = sub.add_parser("set", help="set power and/or fan")
    p.add_argument("--gpus", required=True, help="all or 1,2,5")
    p.add_argument("--power", type=int, default=None)
    p.add_argument("--fan", type=int, default=None)
    p.set_defaults(func=mode_set)

    # restore
    p = sub.add_parser("restore", help="restore auto fan control")
    p.add_argument("--gpus", required=True, help="all or 1,2,5")
    p.set_defaults(func=mode_restore)

    # auto
    p = sub.add_parser("auto", help="auto adjust fan and power based on temperature")
    p.add_argument("--gpus", required=True, help="all or 1,2,5")
    p.add_argument("--time", default=None, help="duration, e.g. 10h, 30m, 3600s. default unlimited")

    p.add_argument("--target-power", type=int, required=True, help="normal target power limit")
    p.add_argument("--min-power", type=int, required=True, help="minimum power limit when hot")
    p.add_argument("--max-power", type=int, required=True, help="maximum allowed power limit")

    p.add_argument("--min-fan", type=int, default=60)
    p.add_argument("--max-fan", type=int, default=95)

    p.add_argument("--hot-temp", type=float, default=78)
    p.add_argument("--cool-temp", type=float, default=68)
    p.add_argument("--critical-temp", type=float, default=84)

    p.add_argument("--fan-step", type=int, default=5)
    p.add_argument("--power-step", type=int, default=25)
    p.add_argument("--interval", type=int, default=15)

    p.add_argument("--restore-on-exit", action="store_true")
    p.set_defaults(func=mode_auto)

    # benchmark
    p = sub.add_parser("benchmark", help="benchmark different power/fan combinations")
    p.add_argument("--gpus", required=True, help="all or 1,2,5")
    p.add_argument("--time", required=True, help="total benchmark duration, e.g. 10h")
    p.add_argument("--stress", required=True, help="stress command")
    p.add_argument("--power-limits", default=",".join(map(str, DEFAULT_POWER_LIMITS)))
    p.add_argument("--fan-speeds", default=",".join(map(str, DEFAULT_FAN_SPEEDS)))
    p.add_argument("--warmup", type=int, default=DEFAULT_WARMUP_SECONDS)
    p.add_argument("--interval", type=int, default=DEFAULT_MONITOR_INTERVAL)
    p.add_argument("--restore-on-exit", action="store_true")
    p.set_defaults(func=mode_benchmark)

    return parser


def main():
    ensure_display()
    parser = build_parser()
    args = parser.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()