PRODIGIT-3366G/prodigit3366g/cli.py

"""CLI tool for controlling the Prodigit 3366G DC Electronic Load via RS-232."""

from __future__ import annotations

import argparse
import csv
import sys
import time

from prodigit3366g.driver import Prodigit3366G


def cmd_identify(load: Prodigit3366G, _args: argparse.Namespace) -> None:
    """Print instrument identity and status."""
    print(f"Model:      {load.name()}")
    print(f"Mode:       {load.get_mode()}")
    print(f"Load:       {'ON' if load.get_load_state() else 'OFF'}")
    print(f"Level:      {load.get_level()}")
    print(f"Sense:      {load.get_sense()}")
    print(f"Dynamic:    {'ON' if load.get_dynamic() else 'OFF'}")
    prot = load.get_protection_status()
    active = [k for k, v in prot.items() if v]
    print(f"Protection: {', '.join(active) if active else 'None'}")


def cmd_measure(load: Prodigit3366G, _args: argparse.Namespace) -> None:
    """Take a single measurement."""
    result = load.measure_all()
    print(f"  Voltage = {result.voltage:>10.4f} V")
    print(f"  Current = {result.current:>10.4f} A")
    print(f"  Power   = {result.power:>10.4f} W")


def cmd_monitor(load: Prodigit3366G, args: argparse.Namespace) -> None:
    """Continuously monitor measurements at an interval."""
    interval = args.interval

    writer = None
    outfile = None
    if args.output:
        outfile = open(args.output, "w", newline="")
        writer = csv.writer(outfile)
        writer.writerow(["timestamp", "voltage_V", "current_A", "power_W"])
        print(f"Logging to {args.output}")

    print(f"{'Time':>10s}  {'Voltage(V)':>12s}  {'Current(A)':>12s}  {'Power(W)':>12s}")
    print("-" * 52)

    try:
        count = 0
        while args.count == 0 or count < args.count:
            result = load.measure_all()
            ts = time.strftime("%H:%M:%S")
            print(
                f"{ts:>10s}  {result.voltage:>12.4f}  "
                f"{result.current:>12.4f}  {result.power:>12.4f}"
            )

            if writer:
                writer.writerow([
                    time.strftime("%Y-%m-%d %H:%M:%S"),
                    f"{result.voltage:.4f}",
                    f"{result.current:.4f}",
                    f"{result.power:.4f}",
                ])
                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(load: Prodigit3366G, args: argparse.Namespace) -> None:
    """Live monitor with real-time graph."""
    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)
    voltages: deque[float] = deque(maxlen=max_points)
    currents: deque[float] = deque(maxlen=max_points)
    powers: deque[float] = deque(maxlen=max_points)
    t0 = time.time()

    writer = None
    outfile = None
    if args.output:
        outfile = open(args.output, "w", newline="")
        writer = csv.writer(outfile)
        writer.writerow(["timestamp", "voltage_V", "current_A", "power_W"])
        print(f"Logging to {args.output}")

    fig, axes = plt.subplots(3, 1, figsize=(12, 9), squeeze=False)
    fig.suptitle("Prodigit 3366G Live Monitor", fontsize=14, fontweight="bold")
    axes = axes.flatten()

    labels = [("Voltage", "V", voltages), ("Current", "A", currents), ("Power", "W", powers)]
    lines = []
    for ax, (title, unit, _) in zip(axes, labels):
        ax.set_ylabel(f"{title} ({unit})")
        ax.set_xlabel("Time (s)")
        ax.grid(True, alpha=0.3)
        ax.set_title(title, fontsize=11)
        line, = ax.plot([], [], linewidth=1.5)
        lines.append(line)
        ax.legend([title], loc="upper left", fontsize=9)

    fig.tight_layout()

    def update(_frame):
        try:
            result = load.measure_all()
        except Exception as e:
            print(f"Read error: {e}")
            return lines

        now = time.time() - t0
        timestamps.append(now)
        voltages.append(result.voltage)
        currents.append(result.current)
        powers.append(result.power)

        ts = time.strftime("%H:%M:%S")
        print(
            f"{ts}  V={result.voltage:.4f}  "
            f"I={result.current:.4f}  P={result.power:.4f}"
        )

        if writer:
            writer.writerow([
                time.strftime("%Y-%m-%d %H:%M:%S"),
                f"{result.voltage:.4f}",
                f"{result.current:.4f}",
                f"{result.power:.4f}",
            ])
            outfile.flush()

        t_list = list(timestamps)
        for line, (_, _, data) in zip(lines, labels):
            line.set_data(t_list, list(data))

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

        return 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_set(load: Prodigit3366G, args: argparse.Namespace) -> None:
    """Set load mode and value."""
    mode = args.mode.upper()
    value = args.value

    load.set_mode(mode)
    if mode == "CC":
        load.set_cc_current(value)
        print(f"Mode: CC, Current: {value:.4f} A")
    elif mode == "CR":
        load.set_cr_resistance(value)
        print(f"Mode: CR, Resistance: {value:.3f} Ohm")
    elif mode == "CV":
        load.set_cv_voltage(value)
        print(f"Mode: CV, Voltage: {value:.4f} V")
    elif mode == "CP":
        load.set_cp_power(value)
        print(f"Mode: CP, Power: {value:.4f} W")


def cmd_load(load: Prodigit3366G, args: argparse.Namespace) -> None:
    """Turn load ON or OFF."""
    if args.state.upper() in ("ON", "1"):
        load.load_on()
        print("Load ON")
    else:
        load.load_off()
        print("Load OFF")


def cmd_battery(load: Prodigit3366G, args: argparse.Namespace) -> None:
    """Battery discharge test control."""
    action = args.action

    if action == "start":
        if args.uvp is not None:
            load.set_battery_test(uvp=args.uvp)
        if args.time is not None:
            load.set_battery_test(time_s=args.time)
        load.battery_test_start()
        print("Battery discharge test started.")
    elif action == "stop":
        load.battery_test_stop()
        results = load.get_battery_results()
        print("Battery test stopped. Results:")
        for k, v in results.items():
            print(f"  {k:>6s} = {v}")
    elif action == "status":
        results = load.get_battery_results()
        print("Battery test results:")
        for k, v in results.items():
            print(f"  {k:>6s} = {v}")


def cmd_send(load: Prodigit3366G, args: argparse.Namespace) -> None:
    """Send a raw command."""
    command = " ".join(args.raw_command)
    result = load.send(command)
    if result is not None:
        print(f"Response: {result}")
    else:
        print("OK")


def cmd_protection(load: Prodigit3366G, args: argparse.Namespace) -> None:
    """Check or clear protection status."""
    if args.clear:
        load.clear_protection()
        print("Protection flags cleared.")
    else:
        prot = load.get_protection_status()
        for name, active in prot.items():
            status = "TRIGGERED" if active else "OK"
            print(f"  {name}: {status}")


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Prodigit 3366G DC Electronic Load RS-232 Tool",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""\
examples:
  %(prog)s identify
  %(prog)s measure
  %(prog)s monitor --interval 1.0 --output data.csv
  %(prog)s live --interval 0.5
  %(prog)s set CC 10.0
  %(prog)s set CV 48.0
  %(prog)s load on
  %(prog)s load off
  %(prog)s battery start --uvp 42.0 --time 3600
  %(prog)s battery stop
  %(prog)s protection
  %(prog)s protection --clear
  %(prog)s send "MODE?"
  %(prog)s send "CC CURR:HIGH 5.0"
""",
    )
    parser.add_argument(
        "-p", "--port", default="COM1",
        help="Serial port (default: COM1)",
    )
    parser.add_argument(
        "-b", "--baudrate", type=int, default=115200,
        help="Baud rate (default: 115200)",
    )
    parser.add_argument(
        "--timeout", type=float, default=2.0,
        help="Serial timeout in seconds (default: 2.0)",
    )

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

    # identify
    sub.add_parser("identify", help="Show instrument identity and status")

    # measure
    sub.add_parser("measure", help="Take a single measurement")

    # monitor
    p_mon = sub.add_parser("monitor", help="Continuously monitor measurements")
    p_mon.add_argument("-i", "--interval", type=float, default=1.0, help="Seconds between readings (default: 1.0)")
    p_mon.add_argument("-n", "--count", type=int, default=0, help="Number of readings (0=infinite)")
    p_mon.add_argument("-o", "--output", help="CSV output file path")

    # live
    p_live = sub.add_parser("live", help="Live monitor with real-time graph")
    p_live.add_argument("-i", "--interval", type=float, default=1.0, help="Seconds between readings (default: 1.0)")
    p_live.add_argument("-o", "--output", help="CSV output file path")
    p_live.add_argument("--history", type=int, default=300, help="Max data points to display (default: 300)")

    # set
    p_set = sub.add_parser("set", help="Set load mode and value")
    p_set.add_argument("mode", choices=["CC", "cc", "CR", "cr", "CV", "cv", "CP", "cp"], help="Load mode")
    p_set.add_argument("value", type=float, help="Set value (A for CC, Ohm for CR, V for CV, W for CP)")

    # load
    p_load = sub.add_parser("load", help="Turn load ON or OFF")
    p_load.add_argument("state", choices=["on", "off", "ON", "OFF", "1", "0"], help="ON or OFF")

    # battery
    p_batt = sub.add_parser("battery", help="Battery discharge test")
    p_batt.add_argument("action", choices=["start", "stop", "status"])
    p_batt.add_argument("--uvp", type=float, help="Under-voltage protection cutoff (V)")
    p_batt.add_argument("--time", type=int, help="Discharge time limit (seconds, 0=unlimited)")

    # protection
    p_prot = sub.add_parser("protection", help="Check or clear protection status")
    p_prot.add_argument("--clear", action="store_true", help="Clear protection flags")

    # send
    p_send = sub.add_parser("send", help="Send raw RS-232 command")
    p_send.add_argument("raw_command", nargs="+", help="Command string (queries auto-detected by '?')")

    args = parser.parse_args()

    dispatch = {
        "identify": cmd_identify,
        "measure": cmd_measure,
        "monitor": cmd_monitor,
        "live": cmd_live,
        "set": cmd_set,
        "load": cmd_load,
        "battery": cmd_battery,
        "protection": cmd_protection,
        "send": cmd_send,
    }

    with Prodigit3366G(port=args.port, baudrate=args.baudrate, timeout=args.timeout) as load:
        load.remote()
        dispatch[args.command](load, args)


if __name__ == "__main__":
    main()