#!/usr/bin/env python3
"""Read-only 50 Hz ball/racket pairing check on the adapter's ZMQ stream."""
import argparse
import json
from pathlib import Path
import sys
import time

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from runtime.rl.config import RLConfig
from runtime.rl.events import BallStream


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--endpoint", default="tcp://127.0.0.1:5556")
    parser.add_argument("--duration-s", type=float, default=10.)
    parser.add_argument("--robot-base-origin-world", type=float, nargs=3, default=(-1.87, 0., -.76))
    args = parser.parse_args()
    stream = BallStream(RLConfig(zmq_endpoint=args.endpoint), args.robot_base_origin_world)
    distances, stamps = [], []
    steps = 0
    start = last_print = time.monotonic()
    try:
        while time.monotonic() - start < args.duration_s:
            stamps.extend(s.source_s for s in stream.drain())
            now = time.monotonic()
            distance = stream.racket_distance(now)
            steps += 1
            if distance is not None:
                distances.append(distance)
            if now - last_print >= 1.:
                with stream.pair_lock:
                    ball = stream.latest_ball
                    racket = stream.racket.samples[-1] if stream.racket.samples else None
                    print(json.dumps(dict(ball_robot_m=ball.position.tolist() if ball else None,
                        racket_robot_m=racket[2].tolist() if racket else None,
                        ball_source_s=ball.source_s if ball else None,
                        racket_source_s=racket[0] if racket else None,
                        distance_m=distance, status="PAIRED" if distance is not None else "NO FRESH PAIR")), flush=True)
                last_print = now
            time.sleep(max(0., .02 - (time.monotonic() - now)))
    finally:
        stream.close()
    print(json.dumps(dict(steps=steps, paired_steps=len(distances), ball_frames=len(stamps),
        ball_source_hz=(len(stamps)-1)/(stamps[-1]-stamps[0]) if len(stamps)>1 else None,
        distance_mean_m=float(np.mean(distances)) if distances else None,
        distance_min_m=min(distances) if distances else None, source_errors=stream.error_generation,
        passed=bool(distances) and stream.error_generation == 0)), flush=True)


if __name__ == "__main__":
    main()