"""CLI for the MPPT Tracker Testbench.

Orchestrates IT6500D (supply) + Prodigit 3366G (load) + HIOKI 3193-10 (meter).
"""

from __future__ import annotations

import argparse
import csv
import sys
import time

from it6500.driver import IT6500
from prodigit3366g.driver import Prodigit3366G
from hioki3193.driver import Hioki3193
from testbench.bench import MPPTTestbench


# ── Instrument connection ─────────────────────────────────────────────


def connect_bench(args: argparse.Namespace) -> MPPTTestbench:
    """Create and return a connected MPPTTestbench from CLI args."""
    supply_addr = MPPTTestbench.find_supply(args.supply_address)
    meter_addr = MPPTTestbench.find_meter(args.meter_address)
    load_port = args.load_port

    print(f"Supply:  {supply_addr}")
    print(f"Load:    {load_port}")
    print(f"Meter:   {meter_addr}")
    print()

    supply = IT6500(supply_addr, timeout_ms=args.timeout)
    load = Prodigit3366G(load_port, baudrate=args.load_baud)
    meter = Hioki3193(meter_addr, timeout_ms=args.timeout)

    return MPPTTestbench(supply, load, meter)


# ── Commands ──────────────────────────────────────────────────────────


def cmd_identify(bench: MPPTTestbench, _args: argparse.Namespace) -> None:
    """Identify all three instruments."""
    print("=== DC Power Supply (IT6500D) ===")
    print(f"  Identity: {bench.supply.idn()}")
    print(f"  Output:   {'ON' if bench.supply.get_output_state() else 'OFF'}")
    print(f"  V set:    {bench.supply.get_voltage():.4f} V")
    print(f"  I set:    {bench.supply.get_current():.4f} A")

    print()
    print("=== DC Electronic Load (Prodigit 3366G) ===")
    print(f"  Model:    {bench.load.name()}")
    print(f"  Load:     {'ON' if bench.load.get_load_state() else 'OFF'}")
    print(f"  Mode:     {bench.load.get_mode()}")

    print()
    print("=== Power Analyzer (HIOKI 3193-10) ===")
    print(f"  Identity: {bench.meter.idn()}")
    print(f"  Options:  {bench.meter.options()}")
    print(f"  Wiring:   {bench.meter.get_wiring_mode()}")


def cmd_setup(bench: MPPTTestbench, _args: argparse.Namespace) -> None:
    """Configure all instruments for MPPT testing."""
    print("Setting up all instruments for MPPT testing...")
    bench.setup_all()
    print()
    print("Supply:  remote mode")
    print("Load:    remote mode")
    print("Meter:   1P2W, DC coupling, auto-range, EFF1=P6/P5")
    print("Display: SELECT16 — U5,I5,P5,EFF1 (left) | U6,I6,P6 (right)")
    print()
    print("Ready. Use 'measure' or 'monitor' to start reading data.")


def cmd_measure(bench: MPPTTestbench, _args: argparse.Namespace) -> None:
    """Take a single measurement from all instruments."""
    data = bench.measure_all()

    print("=== Supply (IT6500D) ===")
    print(f"  Voltage = {data['supply_V']:>10.4f} V")
    print(f"  Current = {data['supply_I']:>10.4f} A")
    print(f"  Power   = {data['supply_P']:>10.4f} W")
    print()
    print("=== Load (Prodigit 3366G) ===")
    print(f"  Voltage = {data['load_V']:>10.4f} V")
    print(f"  Current = {data['load_I']:>10.4f} A")
    print(f"  Power   = {data['load_P']:>10.4f} W")
    print()
    print("=== Meter (HIOKI 3193-10) ===")
    print(f"  Input:  U5={data['meter_U5']:>+12.4E} V  "
          f"I5={data['meter_I5']:>+12.4E} A  "
          f"P5={data['meter_P5']:>+12.4E} W")
    print(f"  Output: U6={data['meter_U6']:>+12.4E} V  "
          f"I6={data['meter_I6']:>+12.4E} A  "
          f"P6={data['meter_P6']:>+12.4E} W")
    print(f"  EFF1  = {data['meter_EFF1']:>+12.4E} %")


def cmd_monitor(bench: MPPTTestbench, args: argparse.Namespace) -> None:
    """Continuously monitor all instruments."""
    interval = args.interval

    writer = None
    outfile = None
    columns = [
        "timestamp",
        "supply_V", "supply_I", "supply_P",
        "load_V", "load_I", "load_P",
        "meter_U5", "meter_I5", "meter_P5",
        "meter_U6", "meter_I6", "meter_P6",
        "meter_EFF1",
    ]
    if args.output:
        outfile = open(args.output, "w", newline="")
        writer = csv.writer(outfile)
        writer.writerow(columns)
        print(f"Logging to {args.output}")

    print(
        f"{'Time':>10s}  "
        f"{'Sup_V':>8s} {'Sup_I':>8s} {'Sup_P':>8s}  "
        f"{'Ld_V':>8s} {'Ld_I':>8s} {'Ld_P':>8s}  "
        f"{'P_in':>10s} {'P_out':>10s} {'EFF%':>8s}"
    )
    print("-" * 105)

    try:
        count = 0
        while args.count == 0 or count < args.count:
            data = bench.measure_all()
            ts = time.strftime("%H:%M:%S")
            print(
                f"{ts:>10s}  "
                f"{data['supply_V']:8.3f} {data['supply_I']:8.3f} {data['supply_P']:8.2f}  "
                f"{data['load_V']:8.3f} {data['load_I']:8.3f} {data['load_P']:8.2f}  "
                f"{data['meter_P5']:>+10.3E} {data['meter_P6']:>+10.3E} {data['meter_EFF1']:8.2f}"
            )

            if writer:
                writer.writerow(
                    [time.strftime("%Y-%m-%d %H:%M:%S")]
                    + [f"{data[c]:.6f}" for c in columns[1:]]
                )
                outfile.flush()

            count += 1
            if args.count == 0 or count < args.count:
                time.sleep(interval)
    except KeyboardInterrupt:
        print("\nMonitoring stopped.")
    finally:
        if outfile:
            outfile.close()
            print(f"Data saved to {args.output}")


def cmd_live(bench: MPPTTestbench, args: argparse.Namespace) -> None:
    """Live monitor with real-time graphs."""
    import matplotlib.pyplot as plt
    from matplotlib.animation import FuncAnimation
    from collections import deque

    max_points = args.history
    interval_ms = int(args.interval * 1000)

    timestamps: deque[float] = deque(maxlen=max_points)
    series: dict[str, deque[float]] = {
        k: deque(maxlen=max_points)
        for k in [
            "supply_P", "load_P",
            "meter_P5", "meter_P6", "meter_EFF1",
            "meter_U5", "meter_U6",
            "meter_I5", "meter_I6",
        ]
    }
    t0 = time.time()

    writer = None
    outfile = None
    columns = [
        "timestamp",
        "supply_V", "supply_I", "supply_P",
        "load_V", "load_I", "load_P",
        "meter_U5", "meter_I5", "meter_P5",
        "meter_U6", "meter_I6", "meter_P6",
        "meter_EFF1",
    ]
    if args.output:
        outfile = open(args.output, "w", newline="")
        writer = csv.writer(outfile)
        writer.writerow(columns)
        print(f"Logging to {args.output}")

    fig, axes = plt.subplots(4, 1, figsize=(14, 12), squeeze=False)
    fig.suptitle("MPPT Testbench Live Monitor", fontsize=14, fontweight="bold")
    axes = axes.flatten()

    # Subplot 0: Power
    axes[0].set_ylabel("Power (W)")
    axes[0].set_title("Power")
    axes[0].grid(True, alpha=0.3)
    line_p5, = axes[0].plot([], [], label="P_in (meter)", linewidth=1.5)
    line_p6, = axes[0].plot([], [], label="P_out (meter)", linewidth=1.5)
    line_sp, = axes[0].plot([], [], label="Supply P", linewidth=1, linestyle="--", alpha=0.6)
    line_lp, = axes[0].plot([], [], label="Load P", linewidth=1, linestyle="--", alpha=0.6)
    axes[0].legend(loc="upper left", fontsize=9)

    # Subplot 1: Efficiency
    axes[1].set_ylabel("Efficiency (%)")
    axes[1].set_title("Efficiency")
    axes[1].grid(True, alpha=0.3)
    line_eff, = axes[1].plot([], [], label="EFF1", linewidth=1.5, color="green")
    axes[1].legend(loc="upper left", fontsize=9)

    # Subplot 2: Voltage
    axes[2].set_ylabel("Voltage (V)")
    axes[2].set_title("Voltage")
    axes[2].grid(True, alpha=0.3)
    line_u5, = axes[2].plot([], [], label="U5 (input)", linewidth=1.5)
    line_u6, = axes[2].plot([], [], label="U6 (output)", linewidth=1.5)
    axes[2].legend(loc="upper left", fontsize=9)

    # Subplot 3: Current
    axes[3].set_ylabel("Current (A)")
    axes[3].set_title("Current")
    axes[3].set_xlabel("Time (s)")
    axes[3].grid(True, alpha=0.3)
    line_i5, = axes[3].plot([], [], label="I5 (input)", linewidth=1.5)
    line_i6, = axes[3].plot([], [], label="I6 (output)", linewidth=1.5)
    axes[3].legend(loc="upper left", fontsize=9)

    fig.tight_layout()

    ERROR_THRESHOLD = 1e90
    all_lines = [line_p5, line_p6, line_sp, line_lp, line_eff, line_u5, line_u6, line_i5, line_i6]

    def _clean(val: float) -> float:
        return float("nan") if abs(val) > ERROR_THRESHOLD else val

    def update(_frame):
        try:
            data = bench.measure_all()
        except Exception as e:
            print(f"Read error: {e}")
            return all_lines

        now = time.time() - t0
        timestamps.append(now)

        for key in series:
            series[key].append(_clean(data[key]))

        ts = time.strftime("%H:%M:%S")
        print(
            f"{ts}  P_in={data['meter_P5']:+.2E}  "
            f"P_out={data['meter_P6']:+.2E}  "
            f"EFF={data['meter_EFF1']:.1f}%"
        )

        if writer:
            writer.writerow(
                [time.strftime("%Y-%m-%d %H:%M:%S")]
                + [f"{data[c]:.6f}" for c in columns[1:]]
            )
            outfile.flush()

        t_list = list(timestamps)
        line_p5.set_data(t_list, list(series["meter_P5"]))
        line_p6.set_data(t_list, list(series["meter_P6"]))
        line_sp.set_data(t_list, list(series["supply_P"]))
        line_lp.set_data(t_list, list(series["load_P"]))
        line_eff.set_data(t_list, list(series["meter_EFF1"]))
        line_u5.set_data(t_list, list(series["meter_U5"]))
        line_u6.set_data(t_list, list(series["meter_U6"]))
        line_i5.set_data(t_list, list(series["meter_I5"]))
        line_i6.set_data(t_list, list(series["meter_I6"]))

        for ax in axes:
            ax.relim()
            ax.autoscale_view()

        return all_lines

    _anim = FuncAnimation(
        fig, update, interval=interval_ms, blit=False, cache_frame_data=False
    )

    try:
        plt.show()
    except KeyboardInterrupt:
        pass
    finally:
        if outfile:
            outfile.close()
            print(f"Data saved to {args.output}")


def cmd_sweep(bench: MPPTTestbench, args: argparse.Namespace) -> None:
    """Run a voltage sweep to characterize MPPT tracking."""
    print(
        f"Voltage sweep: {args.v_start:.1f}V -> {args.v_stop:.1f}V, "
        f"step={args.v_step:.2f}V, I_limit={args.current_limit:.1f}A, "
        f"settle={args.settle:.1f}s"
    )

    # Ensure load is configured
    if args.load_mode:
        bench.load.set_mode(args.load_mode)
    if args.load_value is not None:
        mode = bench.load.get_mode()
        if mode == "CC":
            bench.load.set_cc_current(args.load_value)
        elif mode == "CR":
            bench.load.set_cr_resistance(args.load_value)
        elif mode == "CV":
            bench.load.set_cv_voltage(args.load_value)
        elif mode == "CP":
            bench.load.set_cp_power(args.load_value)
    bench.load.load_on()

    print()
    results = bench.sweep_voltage(
        v_start=args.v_start,
        v_stop=args.v_stop,
        v_step=args.v_step,
        current_limit=args.current_limit,
        settle_time=args.settle,
        load_setpoint=args.load_value if args.load_value is not None else 0.0,
    )

    bench.load.load_off()

    _write_sweep_csv(results, args.output)
    _print_sweep_summary(results)


def _write_sweep_csv(results: list, output: str | None) -> None:
    """Write sweep results to CSV."""
    if not output:
        return
    with open(output, "w", newline="") as f:
        writer = csv.writer(f)
        writer.writerow([
            "voltage_set", "current_limit", "load_setpoint",
            "supply_V", "supply_I", "supply_P",
            "load_V", "load_I", "load_P",
            "input_power", "output_power", "efficiency",
        ])
        for pt in results:
            writer.writerow([
                f"{pt.voltage_set:.4f}",
                f"{pt.current_limit:.4f}",
                f"{pt.load_setpoint:.4f}",
                f"{pt.supply_voltage:.4f}",
                f"{pt.supply_current:.4f}",
                f"{pt.supply_power:.4f}",
                f"{pt.load_voltage:.4f}",
                f"{pt.load_current:.4f}",
                f"{pt.load_power:.4f}",
                f"{pt.input_power:.4f}",
                f"{pt.output_power:.4f}",
                f"{pt.efficiency:.4f}",
            ])
    print(f"\nSweep data saved to {output}")


def _print_sweep_summary(results: list) -> None:
    """Print best-efficiency point from sweep results."""
    if results:
        best = max(results, key=lambda p: p.efficiency)
        print(f"\nBest efficiency: {best.efficiency:.2f}% "
              f"at V_set={best.voltage_set:.2f}V "
              f"I_load={best.load_setpoint:.2f}A "
              f"(P_in={best.input_power:.2f}W, P_out={best.output_power:.2f}W)")


def cmd_sweep_load(bench: MPPTTestbench, args: argparse.Namespace) -> None:
    """Run a load current sweep at a fixed supply voltage."""
    print(
        f"Load current sweep: {args.i_start:.2f}A -> {args.i_stop:.2f}A, "
        f"step={args.i_step:.2f}A, V={args.voltage:.1f}V, "
        f"I_limit={args.current_limit:.1f}A, settle={args.settle:.1f}s"
    )
    print()

    results = bench.sweep_load_current(
        voltage=args.voltage,
        current_limit=args.current_limit,
        i_start=args.i_start,
        i_stop=args.i_stop,
        i_step=args.i_step,
        settle_time=args.settle,
    )

    _write_sweep_csv(results, args.output)
    _print_sweep_summary(results)


def cmd_efficiency(bench: MPPTTestbench, args: argparse.Namespace) -> None:
    """Measure efficiency at a fixed operating point."""
    # Configure load
    if args.load_mode:
        bench.load.set_mode(args.load_mode)
    if args.load_value is not None:
        mode = bench.load.get_mode()
        if mode == "CC":
            bench.load.set_cc_current(args.load_value)
        elif mode == "CR":
            bench.load.set_cr_resistance(args.load_value)
        elif mode == "CV":
            bench.load.set_cv_voltage(args.load_value)
        elif mode == "CP":
            bench.load.set_cp_power(args.load_value)
    bench.load.load_on()

    print(
        f"Measuring efficiency at {args.voltage:.1f}V, "
        f"{args.current_limit:.1f}A limit, "
        f"{args.samples} samples..."
    )
    print()

    result = bench.measure_efficiency(
        voltage=args.voltage,
        current_limit=args.current_limit,
        settle_time=args.settle,
        samples=args.samples,
        sample_interval=args.interval,
    )

    bench.load.load_off()

    print()
    print(f"Average input power:  {result['avg_input_power']:.4f} W")
    print(f"Average output power: {result['avg_output_power']:.4f} W")
    print(f"Average efficiency:   {result['avg_efficiency']:.2f} %")


def cmd_sweep_vi(bench: MPPTTestbench, args: argparse.Namespace) -> None:
    """Run a 2D voltage × load sweep."""
    mode = args.load_mode
    unit = "A" if mode == "CC" else "W"
    print(
        f"2D sweep: V={args.v_start:.1f}-{args.v_stop:.1f}V (step {args.v_step:.1f}), "
        f"{mode}={args.l_start:.2f}-{args.l_stop:.2f}{unit} (step {args.l_step:.2f}), "
        f"I_limit={args.current_limit:.1f}A, settle={args.settle:.1f}s"
    )
    print()

    results = bench.sweep_vi(
        v_start=args.v_start,
        v_stop=args.v_stop,
        v_step=args.v_step,
        l_start=args.l_start,
        l_stop=args.l_stop,
        l_step=args.l_step,
        current_limit=args.current_limit,
        settle_time=args.settle,
        load_mode=mode,
    )

    _write_sweep_csv(results, args.output)
    _print_sweep_summary(results)


def cmd_shade_profile(bench: MPPTTestbench, args: argparse.Namespace) -> None:
    """Run a shade / irradiance profile from a CSV file."""
    steps = MPPTTestbench.load_shade_profile(args.profile)
    duration = steps[-1]["time"] if steps else 0

    print(f"Shade profile: {args.profile}")
    print(f"  {len(steps)} steps over {duration:.0f}s, settle={args.settle:.1f}s")

    # Show what modes/values are used
    modes = {s.get("load_mode", "?") for s in steps}
    voltages = [s["voltage"] for s in steps]
    currents = [s["current_limit"] for s in steps]
    print(
        f"  Voltage: {min(voltages):.1f} - {max(voltages):.1f}V, "
        f"I_limit: {min(currents):.1f} - {max(currents):.1f}A, "
        f"Load modes: {', '.join(sorted(modes))}"
    )
    print()

    results = bench.run_shade_profile(steps, settle_time=args.settle)

    _write_sweep_csv(results, args.output)
    _print_sweep_summary(results)


def cmd_supply(bench: MPPTTestbench, args: argparse.Namespace) -> None:
    """Control the DC supply directly."""
    if args.action == "on":
        bench.supply.output_on()
        print("Supply output ON")
    elif args.action == "off":
        bench.supply.output_off()
        print("Supply output OFF")
    elif args.action == "set":
        if args.voltage is not None and args.current is not None:
            bench.supply.apply(args.voltage, args.current)
            print(f"Supply set: {args.voltage:.4f} V, {args.current:.4f} A")
        elif args.voltage is not None:
            bench.supply.set_voltage(args.voltage)
            print(f"Supply voltage: {args.voltage:.4f} V")
        elif args.current is not None:
            bench.supply.set_current(args.current)
            print(f"Supply current: {args.current:.4f} A")
        else:
            print("Specify --voltage and/or --current")


def cmd_load(bench: MPPTTestbench, args: argparse.Namespace) -> None:
    """Control the electronic load directly."""
    if args.action == "on":
        bench.load.load_on()
        print("Load ON")
    elif args.action == "off":
        bench.load.load_off()
        print("Load OFF")
    elif args.action == "set":
        if args.mode:
            bench.load.set_mode(args.mode)
        if args.value is not None:
            mode = bench.load.get_mode()
            if mode == "CC":
                bench.load.set_cc_current(args.value)
                print(f"Load CC: {args.value:.4f} A")
            elif mode == "CR":
                bench.load.set_cr_resistance(args.value)
                print(f"Load CR: {args.value:.4f} ohm")
            elif mode == "CV":
                bench.load.set_cv_voltage(args.value)
                print(f"Load CV: {args.value:.4f} V")
            elif mode == "CP":
                bench.load.set_cp_power(args.value)
                print(f"Load CP: {args.value:.4f} W")
        elif args.mode:
            print(f"Load mode: {args.mode}")


def cmd_safe_off(bench: MPPTTestbench, _args: argparse.Namespace) -> None:
    """Emergency shutdown: turn off load, then supply."""
    print("Shutting down...")
    bench.safe_off()
    print("  Load OFF")
    print("  Supply OFF")
    print("Done.")


# ── Main ──────────────────────────────────────────────────────────────


def main() -> None:
    parser = argparse.ArgumentParser(
        description="MPPT Tracker Testbench: IT6500D + Prodigit 3366G + HIOKI 3193-10",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""\
examples:
  %(prog)s identify
  %(prog)s setup
  %(prog)s measure
  %(prog)s monitor --interval 1.0 --output data.csv
  %(prog)s live --interval 0.5
  %(prog)s sweep --v-start 10 --v-stop 50 --v-step 1 --current-limit 10 -o sweep.csv
  %(prog)s sweep-load --voltage 75 --current-limit 10 --i-start 1 --i-stop 20 --i-step 1 -o load.csv
  %(prog)s efficiency --voltage 36 --current-limit 10 --samples 10
  %(prog)s sweep-vi --v-start 35 --v-stop 100 --v-step 5 --l-start 0.5 --l-stop 30 --l-step 1 --current-limit 35 -o map.csv
  %(prog)s sweep-vi --v-start 35 --v-stop 100 --v-step 5 --l-start 50 --l-stop 500 --l-step 50 --load-mode CP --current-limit 35 -o map_cp.csv
  %(prog)s shade-profile --profile cloud_pass.csv --settle 2.0 -o shade_results.csv
  %(prog)s supply set --voltage 24 --current 10
  %(prog)s supply on
  %(prog)s load set --mode CC --value 5.0
  %(prog)s load on
  %(prog)s safe-off
""",
    )

    # Global instrument connection args
    parser.add_argument(
        "--supply-address",
        help="IT6500D VISA address (auto-detect if omitted)",
    )
    parser.add_argument(
        "--load-port", default="COM1",
        help="Prodigit 3366G serial port (default: COM1)",
    )
    parser.add_argument(
        "--load-baud", type=int, default=115200,
        help="Prodigit 3366G baud rate (default: 115200)",
    )
    parser.add_argument(
        "--meter-address",
        help="HIOKI 3193-10 VISA address (auto-detect if omitted)",
    )
    parser.add_argument(
        "--timeout", type=int, default=5000,
        help="VISA timeout in ms (default: 5000)",
    )

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

    # identify
    sub.add_parser("identify", help="Identify all connected instruments")

    # setup
    sub.add_parser("setup", help="Configure all instruments for MPPT testing")

    # measure
    sub.add_parser("measure", help="Single measurement from all instruments")

    # monitor
    p_mon = sub.add_parser("monitor", help="Continuous monitoring of all instruments")
    p_mon.add_argument("-i", "--interval", type=float, default=1.0)
    p_mon.add_argument("-n", "--count", type=int, default=0, help="0=infinite")
    p_mon.add_argument("-o", "--output", help="CSV output file")

    # live
    p_live = sub.add_parser("live", help="Live real-time graph of all instruments")
    p_live.add_argument("-i", "--interval", type=float, default=1.0)
    p_live.add_argument("-o", "--output", help="CSV output file")
    p_live.add_argument("--history", type=int, default=300)

    # sweep
    p_sweep = sub.add_parser("sweep", help="Voltage sweep to characterize MPPT tracking")
    p_sweep.add_argument("--v-start", type=float, required=True, help="Start voltage (V)")
    p_sweep.add_argument("--v-stop", type=float, required=True, help="Stop voltage (V)")
    p_sweep.add_argument("--v-step", type=float, required=True, help="Voltage step (V)")
    p_sweep.add_argument("--current-limit", type=float, required=True, help="Current limit (A)")
    p_sweep.add_argument("--settle", type=float, default=1.0, help="Settle time per step (s)")
    p_sweep.add_argument("--load-mode", choices=["CC", "CR", "CV", "CP"], help="Set load mode before sweep")
    p_sweep.add_argument("--load-value", type=float, help="Set load value before sweep")
    p_sweep.add_argument("-o", "--output", help="CSV output file")

    # sweep-load
    p_swl = sub.add_parser("sweep-load", help="Load current sweep at fixed supply voltage")
    p_swl.add_argument("--voltage", type=float, required=True, help="Fixed supply voltage (V)")
    p_swl.add_argument("--current-limit", type=float, required=True, help="Supply current limit (A)")
    p_swl.add_argument("--i-start", type=float, required=True, help="Start load current (A)")
    p_swl.add_argument("--i-stop", type=float, required=True, help="Stop load current (A)")
    p_swl.add_argument("--i-step", type=float, required=True, help="Current step (A)")
    p_swl.add_argument("--settle", type=float, default=1.0, help="Settle time per step (s)")
    p_swl.add_argument("-o", "--output", help="CSV output file")

    # efficiency
    p_eff = sub.add_parser("efficiency", help="Measure efficiency at fixed operating point")
    p_eff.add_argument("--voltage", type=float, required=True, help="Supply voltage (V)")
    p_eff.add_argument("--current-limit", type=float, required=True, help="Current limit (A)")
    p_eff.add_argument("--samples", type=int, default=5, help="Number of readings to average")
    p_eff.add_argument("--settle", type=float, default=2.0, help="Initial settle time (s)")
    p_eff.add_argument("-i", "--interval", type=float, default=1.0, help="Interval between samples")
    p_eff.add_argument("--load-mode", choices=["CC", "CR", "CV", "CP"])
    p_eff.add_argument("--load-value", type=float)

    # sweep-vi (2D)
    p_svi = sub.add_parser("sweep-vi", help="2D voltage × load sweep (efficiency map)")
    p_svi.add_argument("--v-start", type=float, required=True, help="Start voltage (V)")
    p_svi.add_argument("--v-stop", type=float, required=True, help="Stop voltage (V)")
    p_svi.add_argument("--v-step", type=float, required=True, help="Voltage step (V)")
    p_svi.add_argument("--l-start", type=float, required=True, help="Start load setpoint (A for CC, W for CP)")
    p_svi.add_argument("--l-stop", type=float, required=True, help="Stop load setpoint")
    p_svi.add_argument("--l-step", type=float, required=True, help="Load step size")
    p_svi.add_argument("--load-mode", choices=["CC", "CP"], default="CC", help="Load mode: CC (current) or CP (power)")
    p_svi.add_argument("--current-limit", type=float, required=True, help="Supply current limit (A)")
    p_svi.add_argument("--settle", type=float, default=2.0, help="Settle time per step (s)")
    p_svi.add_argument("-o", "--output", help="CSV output file")

    # shade-profile
    p_shade = sub.add_parser("shade-profile", help="Run a shade/irradiance profile from CSV")
    p_shade.add_argument("--profile", required=True, help="Profile CSV file (time,voltage,current_limit,...)")
    p_shade.add_argument("--settle", type=float, default=2.0, help="Settle time per step (s)")
    p_shade.add_argument("-o", "--output", help="CSV output file for results")

    # supply (direct control)
    p_sup = sub.add_parser("supply", help="Direct supply control")
    p_sup_sub = p_sup.add_subparsers(dest="action", required=True)
    p_sup_sub.add_parser("on", help="Turn supply output ON")
    p_sup_sub.add_parser("off", help="Turn supply output OFF")
    p_sup_set = p_sup_sub.add_parser("set", help="Set supply voltage/current")
    p_sup_set.add_argument("-v", "--voltage", type=float)
    p_sup_set.add_argument("-c", "--current", type=float)

    # load (direct control)
    p_ld = sub.add_parser("load", help="Direct load control")
    p_ld_sub = p_ld.add_subparsers(dest="action", required=True)
    p_ld_sub.add_parser("on", help="Turn load ON")
    p_ld_sub.add_parser("off", help="Turn load OFF")
    p_ld_set = p_ld_sub.add_parser("set", help="Set load mode and value")
    p_ld_set.add_argument("-m", "--mode", choices=["CC", "CR", "CV", "CP"])
    p_ld_set.add_argument("-v", "--value", type=float, help="Setpoint value")

    # safe-off
    sub.add_parser("safe-off", help="Emergency: turn off load and supply")

    args = parser.parse_args()

    dispatch = {
        "identify": cmd_identify,
        "setup": cmd_setup,
        "measure": cmd_measure,
        "monitor": cmd_monitor,
        "live": cmd_live,
        "sweep": cmd_sweep,
        "sweep-load": cmd_sweep_load,
        "efficiency": cmd_efficiency,
        "sweep-vi": cmd_sweep_vi,
        "shade-profile": cmd_shade_profile,
        "supply": cmd_supply,
        "load": cmd_load,
        "safe-off": cmd_safe_off,
    }

    bench = connect_bench(args)
    try:
        bench.supply.remote()
        bench.load.remote()
        dispatch[args.command](bench, args)
    except KeyboardInterrupt:
        print("\nInterrupted.")
    finally:
        bench.close()


if __name__ == "__main__":
    main()