#!/usr/bin/env python
"""Where do the 1–2 seconds between "decision locked" and "drone moves" actually go?

Motivation: the working 4–5 s SSVEP loop locks a command, and the Tello only visibly reacts
1–2 s later. That extra delay is as expensive as another SSVEP epoch, so it is worth pinning
down instead of guessing. There is no documented "wait before executing" in the Tello firmware
— the SDK 2.0 guide specifies no command queue and no pre-execution hold — so the delay has to
come from somewhere in this stack:

    A. host software      waiting for the previous command's "ok"  (blocking semantics)
    B. Wi-Fi link         RTT, plus contention if the video stream is on
    C. distance commands  `forward x` is x=20..500 cm at `speed` cm/s → travel time is 1-2 s
                          by arithmetic alone, before any accel/brake ramp
    D. flight controller  accel ramp, brake, VPS re-lock after the move

This script separates them. It deliberately uses **raw UDP sockets, not djitellopy**, so the
library is not part of the measurement; if raw sockets are fast and your djitellopy loop is
slow, the answer is (A).

Onset is measured from the state broadcast on port 8890 (~10 Hz): timestamp the moment the
command leaves the socket, then find the first state packet whose velocity leaves the hover
baseline. The threshold is derived from the hover noise of that same trial, so it does not
depend on the (undocumented) units of vgx/vgy/vgz. Resolution is one state packet ≈ 100 ms —
plenty to tell 80 ms apart from 1200 ms, but do not read the third digit.

  python src/control/latency_test.py --link            # SAFE: no flight, no motors
  python src/control/latency_test.py --flight --yes    # FLIES. Needs ~3 m of clear space.

`--link` alone already tests the cheapest hypothesis (C): it prints `speed?`. If the drone is
set to e.g. 30 cm/s, then `forward 30` takes 1.0 s of pure travel and nothing is broken.
"""
from __future__ import annotations

import argparse
import socket
import statistics
import threading
import time
from pathlib import Path

TELLO = ("192.168.10.1", 8889)
CMD_PORT = 9010          # local port for command/response
STATE_PORT = 8890        # fixed by the SDK: the drone broadcasts state here

ROOT = Path(__file__).resolve().parents[2]


# --------------------------------------------------------------------------- state listener
class StateStream:
    """Background reader for the port-8890 telemetry broadcast, with host timestamps."""

    def __init__(self):
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.sock.bind(("", STATE_PORT))
        self.sock.settimeout(0.5)
        self.samples: list[tuple[float, dict]] = []      # (t_host, fields)
        self._run = False
        self._th = None

    def start(self):
        self._run = True
        self._th = threading.Thread(target=self._loop, daemon=True)
        self._th.start()

    def _loop(self):
        while self._run:
            try:
                pkt = self.sock.recv(2048).decode(errors="replace")
            except socket.timeout:
                continue
            except OSError:
                break
            t = time.perf_counter()
            d = {}
            for kv in pkt.strip().rstrip(";").split(";"):
                if ":" in kv:
                    k, _, v = kv.partition(":")
                    try:
                        d[k] = float(v)
                    except ValueError:
                        pass
            if d:
                self.samples.append((t, d))

    def stop(self):
        self._run = False
        if self._th:
            self._th.join(timeout=1.0)
        self.sock.close()

    def since(self, t0):
        return [(t, d) for t, d in self.samples if t >= t0]

    def window(self, t0, t1):
        return [(t, d) for t, d in self.samples if t0 <= t <= t1]


def speed_mag(d):
    """Horizontal+vertical speed magnitude, unit-agnostic (we only compare against baseline)."""
    return max(abs(d.get("vgx", 0.0)), abs(d.get("vgy", 0.0)), abs(d.get("vgz", 0.0)))


# --------------------------------------------------------------------------- command channel
class Link:
    def __init__(self, verbose=True):
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.sock.bind(("", CMD_PORT))
        self.verbose = verbose

    def send(self, cmd, timeout=10.0):
        """Send and wait for the reply. Returns (t_sent, t_reply, text). t_reply is None on timeout.

        Blocking on purpose: the Tello does not acknowledge a motion command until the motion
        has FINISHED, so (t_reply - t_sent) is the command's completion time, not its latency."""
        self.sock.settimeout(timeout)
        t0 = time.perf_counter()
        self.sock.sendto(cmd.encode(), TELLO)
        try:
            txt = self.sock.recv(1024).decode(errors="replace").strip()
            t1 = time.perf_counter()
        except socket.timeout:
            txt, t1 = None, None
        if self.verbose:
            dt = "  --  " if t1 is None else f"{(t1 - t0) * 1e3:7.1f} ms"
            print(f"    {cmd:<20s} → {str(txt):<12s} {dt}")
        return t0, t1, txt

    def fire(self, cmd):
        """Send without waiting — this is how `rc` is meant to be used."""
        t0 = time.perf_counter()
        self.sock.sendto(cmd.encode(), TELLO)
        return t0

    def close(self):
        self.sock.close()


# --------------------------------------------------------------------------- link-only test
def test_link(args):
    print("\n=== 链路测试(不起飞,不通电马达)===")
    link = Link()
    state = StateStream()
    state.start()

    t0, t1, txt = link.send("command", timeout=5.0)
    if txt is None:
        print("\n  没有响应。检查 Mac 是否连在 TELLO-XXXXXX 这个 WiFi 上(应拿到 192.168.10.x)。")
        state.stop(); link.close(); return

    print("\n  -- 只读查询(这些不动马达)--")
    info = {}
    for q in ("sdk?", "sn?", "battery?", "speed?", "wifi?", "time?"):
        _, _, r = link.send(q, timeout=3.0)
        info[q] = r

    print("\n  -- 空指令往返时延 × %d --" % args.pings)
    rtt = []
    for _ in range(args.pings):
        a, b, r = link.send("battery?", timeout=3.0)
        if b is not None:
            rtt.append((b - a) * 1e3)
        time.sleep(0.05)

    print("\n  -- 状态广播(8890)节奏 --")
    time.sleep(2.0)
    ts = [t for t, _ in state.samples]
    gaps = [(b - a) * 1e3 for a, b in zip(ts, ts[1:])]

    print("\n--- 结果 ---")
    if rtt:
        print(f"  指令往返 RTT   中位 {statistics.median(rtt):.1f} ms   "
              f"最小 {min(rtt):.1f}   最大 {max(rtt):.1f} ms")
    if gaps:
        print(f"  状态包间隔     中位 {statistics.median(gaps):.1f} ms  "
              f"→ 约 {1000/statistics.median(gaps):.1f} Hz  (共 {len(ts)} 包)")
    else:
        print("  状态包         一个都没收到 —— 8890 端口可能被别的程序占了")

    sp = info.get("speed?")
    print(f"\n  当前 speed 设置: {sp}")
    try:
        v = float(str(sp).strip())
        for d in (20, 30, 50):
            print(f"    → `forward {d}` 纯行程时间下限 = {d}/{v:.0f} = {d/v:.2f} s"
                  f"{'   ← 这就够解释 1-2 秒了' if d / v > 0.8 else ''}")
        if v < 80:
            print(f"\n  ** speed={v:.0f} cm/s 偏低。发一次 `speed 100` 就能直接砍掉大部分延迟。 **")
    except (TypeError, ValueError):
        pass

    if rtt and statistics.median(rtt) < 100:
        print("\n  链路本身很快 —— 1-2 秒不是 WiFi 的问题,继续跑 --flight 看运动起始时刻。")
    state.stop(); link.close()


# --------------------------------------------------------------------------- flight test
def onset_and_end(state, t_send, baseline_v, t_limit):
    """First state packet after t_send whose speed leaves the hover baseline, and when it settles."""
    thr = max(baseline_v * 3.0, 5.0)
    seq = state.window(t_send, t_send + t_limit)
    t_on = t_end = None
    for t, d in seq:
        v = speed_mag(d)
        if t_on is None and v > thr:
            t_on = t
        elif t_on is not None and v <= thr:
            t_end = t
            break
    return t_on, t_end, thr


def test_flight(args):
    print("\n=== 飞行测试 ===")
    print("  会起飞并前后小幅移动。请确认周围 3 米内无人无障碍物,电池 > 40%。")
    if not args.yes:
        print("  加 --yes 才会真的起飞。已退出。")
        return

    link = Link()
    state = StateStream()
    state.start()
    _, _, txt = link.send("command", timeout=5.0)
    if txt is None:
        print("  连不上,退出。"); state.stop(); link.close(); return
    link.send("battery?", timeout=3.0)
    link.send(f"speed {args.speed}", timeout=3.0)

    for i in range(args.countdown, 0, -1):
        print(f"    {i} …"); time.sleep(1.0)

    rows = []
    try:
        print("\n  起飞 …")
        t0, t1, r = link.send("takeoff", timeout=25.0)
        if r is None or "ok" not in str(r).lower():
            print("  起飞失败,退出。"); return
        rows.append(("takeoff", t1 - t0, None, None))
        time.sleep(args.settle)

        for k in range(args.trials):
            direction = "forward" if k % 2 == 0 else "back"

            # ---- (1) distance command: `forward 30`
            base = [speed_mag(d) for _, d in state.window(time.perf_counter() - 0.8,
                                                          time.perf_counter())]
            bv = statistics.median(base) if base else 0.0
            ts, tr, r = link.send(f"{direction} {args.dist}", timeout=15.0)
            t_on, t_end, thr = onset_and_end(state, ts, bv, 12.0)
            rows.append((f"{direction} {args.dist}",
                         None if tr is None else tr - ts,
                         None if t_on is None else t_on - ts,
                         None if t_end is None else t_end - ts))
            print(f"      起始 {'--' if t_on is None else f'{(t_on-ts)*1e3:6.0f} ms'}   "
                  f"运动结束 {'--' if t_end is None else f'{(t_end-ts)*1e3:6.0f} ms'}   "
                  f"(阈值 {thr:.1f})")
            time.sleep(args.settle)

            # ---- (2) rc setpoint: fire-and-forget, no ack to wait for
            rev = -1 if direction == "back" else 1
            base = [speed_mag(d) for _, d in state.window(time.perf_counter() - 0.8,
                                                          time.perf_counter())]
            bv = statistics.median(base) if base else 0.0
            ts = link.fire(f"rc 0 {rev * args.rc} 0 0")
            deadline = ts + args.rc_hold
            while time.perf_counter() < deadline:
                time.sleep(0.02)
            link.fire("rc 0 0 0 0")
            t_on, _, thr = onset_and_end(state, ts, bv, 3.0)
            rows.append((f"rc {rev * args.rc}", None,
                         None if t_on is None else t_on - ts, None))
            print(f"    rc {rev*args.rc:<4d}          "
                  f"起始 {'--' if t_on is None else f'{(t_on-ts)*1e3:6.0f} ms'}"
                  f"   (无 ack,不阻塞)")
            time.sleep(args.settle)
    finally:
        print("\n  降落 …")
        link.fire("rc 0 0 0 0")
        link.send("land", timeout=20.0)
        time.sleep(1.0)
        state.stop()
        link.close()

    summarise(rows)
    plot(rows, args)


def summarise(rows):
    print("\n--- 汇总(ms)---")
    print(f"  {'命令':<16s}{'到 ack':>10s}{'运动起始':>10s}{'运动结束':>10s}")
    for name, ack, on, end in rows:
        f = lambda x: "   --" if x is None else f"{x*1e3:7.0f}"
        print(f"  {name:<16s}{f(ack):>10s}{f(on):>10s}{f(end):>10s}")

    dist_on = [on for n, _, on, _ in rows if on is not None and not n.startswith(("rc", "takeoff"))]
    rc_on = [on for n, _, on, _ in rows if on is not None and n.startswith("rc")]
    dist_ack = [a for n, a, _, _ in rows if a is not None and not n.startswith(("rc", "takeoff"))]
    print()
    if dist_on:
        print(f"  距离指令 起始延迟 中位 {statistics.median(dist_on)*1e3:.0f} ms")
    if dist_ack:
        print(f"  距离指令 到 ack  中位 {statistics.median(dist_ack)*1e3:.0f} ms"
              f"   ← 阻塞式调用会等这么久")
    if rc_on:
        print(f"  rc 指令  起始延迟 中位 {statistics.median(rc_on)*1e3:.0f} ms")
    if dist_on and rc_on:
        a, b = statistics.median(dist_on), statistics.median(rc_on)
        print("\n  结论:", "rc 明显更快 → 换成 rc 速度设定点即可" if a - b > 0.2
              else "两者起始延迟接近 → 瓶颈在链路或飞控,不在指令类型")


def plot(rows, args):
    try:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
    except ImportError:
        print("  (没装 matplotlib,跳过出图)"); return

    names = [r[0] for r in rows]
    ack = [0 if r[1] is None else r[1] * 1e3 for r in rows]
    on = [0 if r[2] is None else r[2] * 1e3 for r in rows]
    y = range(len(rows))

    fig, ax = plt.subplots(figsize=(9, 0.45 * len(rows) + 2.2))
    ax.barh(list(y), ack, color="#c9d6e4", label="到 ack(指令完成)")
    ax.barh(list(y), on, height=0.45, color="#c0392b", label="运动起始")
    ax.set_yticks(list(y)); ax.set_yticklabels(names, fontsize=8)
    ax.invert_yaxis()
    ax.set_xlabel("延迟 (ms)")
    ax.set_title("Tello 指令延迟:运动何时开始 vs ack 何时返回")
    ax.legend(fontsize=8); ax.grid(axis="x", alpha=0.3)
    fig.tight_layout()
    out = ROOT / "results" / "tello_latency.png"
    out.parent.mkdir(exist_ok=True)
    fig.savefig(out, dpi=140)
    print(f"\n  图: {out}")


def main():
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--link", action="store_true", help="只测链路,不起飞(安全)")
    p.add_argument("--flight", action="store_true", help="起飞实测运动起始延迟")
    p.add_argument("--yes", action="store_true", help="确认可以起飞")
    p.add_argument("--pings", type=int, default=20)
    p.add_argument("--trials", type=int, default=4)
    p.add_argument("--dist", type=int, default=30, help="距离指令的 cm(SDK 最小 20)")
    p.add_argument("--speed", type=int, default=100, help="测试前设定的 speed cm/s")
    p.add_argument("--rc", type=int, default=60, help="rc 速度设定点 (-100..100)")
    p.add_argument("--rc-hold", type=float, default=0.6, help="rc 保持多久后归零")
    p.add_argument("--settle", type=float, default=2.5, help="每次动作之间的稳定等待")
    p.add_argument("--countdown", type=int, default=5)
    a = p.parse_args()

    if not (a.link or a.flight):
        a.link = True
    if a.link:
        test_link(a)
    if a.flight:
        test_flight(a)


if __name__ == "__main__":
    main()