Hopi/Hopi.py

import argparse
import time
from datetime import datetime, timezone

from hopi.data_handler import PowerMeterDataHandler
from hopi.gui_plotter import PlotConfig, RealtimePlotter
from hopi.mqtt_client import MqttConfig, MqttJsonSubscriber
from hopi.serial_client import ModbusRtuSerialClient, SerialConfig


SLAVE_ID = 1
START_REG = 0
NUM_WORDS = 20
READ_MAX_BYTES = 5 + (2 * NUM_WORDS)


def build_arg_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Read and print power meter values over Modbus RTU.")
    parser.add_argument(
        "--source",
        choices=["usb", "mqtt"],
        default="usb",
        help="Data source: 'usb' (Modbus RTU over serial) or 'mqtt' (JSON over MQTT)",
    )

    parser.add_argument("--port", default="/dev/ttyUSB0")
    parser.add_argument("--baud", type=int, default=9600)
    parser.add_argument("--timeout", type=float, default=0.1)

    parser.add_argument("--mqtt-server", default="localhost:1883", help="MQTT broker host or host:port")
    parser.add_argument("--mqtt-topic", default="pico-usb-host-modbus/modbus/readings", help="MQTT topic to subscribe to")
    parser.add_argument("--mqtt-username", default=None)
    parser.add_argument("--mqtt-password", default=None)
    parser.add_argument("--mqtt-client-id", default=None)

    parser.add_argument("--daemon", action="store_true", help="Run continuously until stopped (Ctrl+C)")
    parser.add_argument("--csv", default="hopi_readings.csv", help="CSV output path for daemon mode")
    parser.add_argument("--interval", type=float, default=1.0, help="Seconds between reads in daemon mode")
    parser.add_argument("--gui", action="store_true", help="Show realtime plot in daemon mode")
    parser.add_argument("--max-points", type=int, default=300, help="Max points to keep in the plot")
    return parser


def main() -> None:
    args = build_arg_parser().parse_args()

    handler = PowerMeterDataHandler()

    if args.source == "usb":
        config = SerialConfig(port=args.port, baudrate=args.baud, timeout=args.timeout)
        client = ModbusRtuSerialClient(config)
    else:
        client = None

    if args.daemon:
        if args.timeout <= 0:
            raise ValueError("--timeout must be > 0")
        if args.interval <= 0:
            raise ValueError("--interval must be > 0")
        if args.source == "usb" and args.interval <= args.timeout:
            raise ValueError("--interval must be > --timeout")

        t0 = time.monotonic()
        next_deadline = t0

        plotter = None
        if args.gui:
            series_names = handler.flat_readings_keys()  # stable order
            plotter = RealtimePlotter(series_names, config=PlotConfig(max_points=max(10, args.max_points)))
            plotter.open()

        if args.source == "usb":
            assert client is not None
            with client, handler.open_csv_logger(args.csv) as logger:
                try:
                    while True:
                        if plotter is not None and not plotter.is_open():
                            return

                        response = client.read_holding_registers(
                            slave_id=SLAVE_ID,
                            start_register=START_REG,
                            register_count=NUM_WORDS,
                            read_max_bytes=READ_MAX_BYTES,
                            exact_response_length=True,
                        )
                        registers = handler.parse_read_holding_registers_response(response)
                        readings = handler.decode_readings(registers)

                        row = {"timestamp": datetime.now(timezone.utc).isoformat()}
                        flat = handler.readings_to_flat_dict(readings)
                        row.update(flat)
                        logger.log(row)

                        if plotter is not None:
                            plotter.update(time.monotonic() - t0, flat)
                            plotter.process_events()

                        next_deadline += args.interval
                        sleep_for = next_deadline - time.monotonic()
                        if sleep_for > 0:
                            time.sleep(sleep_for)
                except KeyboardInterrupt:
                    return

        mqtt_cfg = MqttConfig(
            server=args.mqtt_server,
            topic=args.mqtt_topic,
            username=args.mqtt_username,
            password=args.mqtt_password,
            client_id=args.mqtt_client_id,
        )

        with MqttJsonSubscriber(mqtt_cfg) as sub, handler.open_csv_logger(args.csv) as logger:
            try:
                while True:
                    if plotter is not None and not plotter.is_open():
                        return

                    payload = sub.get(timeout=args.interval)
                    if payload is None:
                        if plotter is not None:
                            plotter.process_events()
                        continue

                    readings = handler.decode_readings_from_mqtt_payload(payload)
                    row = {"timestamp": datetime.now(timezone.utc).isoformat()}
                    flat = handler.readings_to_flat_dict(readings)
                    row.update(flat)
                    logger.log(row)

                    if plotter is not None:
                        plotter.update(time.monotonic() - t0, flat)
                        plotter.process_events()
            except KeyboardInterrupt:
                return

    if args.source == "mqtt":
        mqtt_cfg = MqttConfig(
            server=args.mqtt_server,
            topic=args.mqtt_topic,
            username=args.mqtt_username,
            password=args.mqtt_password,
            client_id=args.mqtt_client_id,
        )
        with MqttJsonSubscriber(mqtt_cfg) as sub:
            payload = sub.get(timeout=max(15.0, args.timeout))
            if payload is None:
                raise RuntimeError("No MQTT message received (timeout)")
            readings = handler.decode_readings_from_mqtt_payload(payload)
            handler.print_readings_json(readings)
        return

    assert client is not None
    with client:
        response = client.read_holding_registers(
            slave_id=SLAVE_ID,
            start_register=START_REG,
            register_count=NUM_WORDS,
            read_max_bytes=READ_MAX_BYTES,
            exact_response_length=True,
        )

    registers = handler.parse_read_holding_registers_response(response)
    readings = handler.decode_readings(registers)
    handler.print_readings_json(readings)


if __name__ == "__main__":
    main()