#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""jog_probe_pi.py — 根元 j0 で 4 つの動かし方を測り, mechArm がどこまでなめらかに動くかを見る (2026-11-03)

Mac の tools/jog_probe.sh から呼ぶ (デーモンを止めてから, シリアルに直接つなぐ)。
Pi の python3 は 3.6, pymycobot 2.9.6 (PYTHONPATH=/usr/local/lib/python2.7/dist-packages)。

測るもの (どれも j0 だけを動かす. 始める前に HOME 近くにいること. 各テストのあと始めの角度へ戻す):
  0. 読み取り: get_angles を 40 回. 1 回の時間と, 空で返る割合
  A. 1 回の send_angles (+10°, speed 20): 着くまでの時間と最高速度 (アーム自体の速さ)
  B. 送り直し (今のゲームパッドの長押しと同じ): 0.1 秒おきに「読んだ角度 + 3°」を send_angles. 1 秒
  C. B と同じ. ただし set_fresh_mode(1) (前の動きを待たずに上書き)
  D. jog_angle(1, 向き, speed) を 1 秒 → jog_stop. speed 10 と 30
各テストで: 進んだ角度・平均の速さ・0.1 秒ごとの速さのばらつき (なめらかさ)・止まった区間の割合・
止める命令のあとに進んだ角度 (行きすぎ) を出す。生の記録 (時刻と角度) は --out の JSON に入る。

守り: j0 が始めの角度から JOG_GUARD_DEG (25°) 離れたらそのテストを止める (jog_stop・stop)。
q と Enter (または Ctrl-C・ssh が切れたとき) で止める。終わるとき (止めたときも) jog_stop と set_fresh_mode(0)。
send_coords() は使わない。release 系は呼ばない。
"""

from __future__ import print_function

import argparse
import json
import select
import signal
import statistics
import sys
import time

SERIAL_PORT = "/dev/ttyAMA0"
BAUDRATE = 1000000
HOME = [0, 0, 0, 0, 0, -30]
SPEED = 20
HOME_TOL_DEG = 5.0
JOG_GUARD_DEG = 25.0
STEP_A_DEG = 10.0
STREAM_SEC = 1.0
STREAM_INTERVAL_SEC = 0.1
STREAM_AHEAD_DEG = 3.0
JOG_SEC = 1.0
JOG_SPEEDS = (10, 30)
AFTER_STOP_SEC = 1.0
RETURN_TIMEOUT_SEC = 10.0
BIN_SEC = 0.1
GET_ANGLES_RETRIES = 12


class Quit(Exception):
    pass


def now():
    return time.time()


def stdin_quit():
    r, _w, _x = select.select([sys.stdin], [], [], 0)
    if r:
        line = sys.stdin.readline()
        if line.strip().lower() == "q":
            raise Quit()


def ask(msg):
    print(msg + "  (Enter で進む / q と Enter で終わる)", flush=True)
    line = sys.stdin.readline()
    if not line or line.strip().lower() == "q":
        raise Quit()


def read_angles(mc, tries=GET_ANGLES_RETRIES):
    for _ in range(tries):
        a = mc.get_angles()
        if a and len(a) == 6:
            return [float(x) for x in a]
        time.sleep(0.02)
    return None


def wait_arrive(mc, target, timeout):
    t0 = now()
    a = None
    while now() - t0 < timeout:
        stdin_quit()
        a = read_angles(mc)
        if a is not None and max(abs(x - y) for x, y in zip(a, target)) <= 2.0:
            return a, True
        time.sleep(0.1)
    return a, False


def summarize(samples, t_cmd0, t_stop, sign):
    """samples: [(t, j0)]. t_cmd0 = 動かし始めた時刻, t_stop = 止める命令 (なければ None)."""
    act = [(t, q) for t, q in samples if t >= t_cmd0 and (t_stop is None or t <= t_stop)]
    out = {"n_samples": len(samples)}
    if len(act) < 2:
        return out
    travel = (act[-1][1] - act[0][1]) * sign
    dur = act[-1][0] - act[0][0]
    out.update(travel_deg=round(travel, 2), active_sec=round(dur, 2),
               mean_speed=round(travel / dur, 2) if dur > 0 else None)
    # 0.1 秒ごとの速さ (最初の 0.2 秒は加速として除く)
    bins = []
    t = act[0][0] + 0.2
    while t + BIN_SEC <= act[-1][0]:
        seg = [(ts, q) for ts, q in act if t <= ts <= t + BIN_SEC]
        if len(seg) >= 2:
            bins.append((seg[-1][1] - seg[0][1]) * sign / (seg[-1][0] - seg[0][0]) if seg[-1][0] > seg[0][0] else 0.0)
        t += BIN_SEC
    if bins:
        med = statistics.median(bins)
        out.update(bin_speeds=[round(b, 1) for b in bins], median_speed=round(med, 2),
                   max_speed=round(max(bins), 2),
                   speed_cv=round(statistics.pstdev(bins) / med, 2) if med > 0.5 else None,
                   stall_ratio=round(sum(1 for b in bins if b < 0.25 * med) / len(bins), 2) if med > 0.5 else None)
    if t_stop is not None:
        before = [q for ts, q in samples if ts <= t_stop]
        after = [(ts, q) for ts, q in samples if ts > t_stop]
        if before and after:
            q_stop = before[-1]
            out["overshoot_deg"] = round((after[-1][1] - q_stop) * sign, 2)
            moving = [ts for (ts, q), (ts2, q2) in zip(after, after[1:]) if abs(q2 - q) > 0.1]
            out["stop_latency_sec"] = round(moving[-1] - t_stop, 2) if moving else 0.0
    return out


def run_sampling(mc, start, sec, sign, on_tick=None, label=""):
    """sec 秒, できるだけ速く j0 を読む. on_tick(t, angles) を毎回呼ぶ. 守りを超えたら止める."""
    samples = []
    t0 = now()
    while now() - t0 < sec:
        stdin_quit()
        a = mc.get_angles()
        t = now()
        if a and len(a) == 6:
            samples.append((t, float(a[0])))
            if abs(float(a[0]) - start[0]) > JOG_GUARD_DEG:
                print("  守り: j0 が始めから {:.1f}° 離れたので止めます".format(float(a[0]) - start[0]), flush=True)
                return samples, True
            if on_tick is not None:
                on_tick(t, [float(x) for x in a])
    return samples, False


def back_to_start(mc, start):
    mc.send_angles(start, SPEED)
    a, ok = wait_arrive(mc, start, RETURN_TIMEOUT_SEC)
    print("  始めの角度へ戻した: " + ("着いた" if ok else "着かない") + " j0={}".format(None if a is None else round(a[0], 2)), flush=True)
    time.sleep(0.5)


def test_read(mc):
    lat, empty = [], 0
    for _ in range(40):
        t = now()
        a = mc.get_angles()
        lat.append(now() - t)
        if not (a and len(a) == 6):
            empty += 1
    r = {"median_ms": round(statistics.median(lat) * 1000, 1), "max_ms": round(max(lat) * 1000, 1),
         "empty_ratio": round(empty / 40.0, 2)}
    print("  読み取り 1 回: 中央 {median_ms} ms / 最大 {max_ms} ms / 空で返った割合 {empty_ratio}".format(**r), flush=True)
    return r


def test_step(mc, start):
    target = list(start)
    target[0] = start[0] + STEP_A_DEG
    t_cmd = now()
    mc.send_angles(target, SPEED)
    samples, guard = run_sampling(mc, start, 2.0, +1)
    s = summarize(samples, t_cmd, None, +1)
    reached = [t for t, q in samples if abs(q - target[0]) <= 1.0]
    s["arrive_sec"] = round(reached[0] - t_cmd, 2) if reached else None
    return s, samples, t_cmd, None


def test_stream(mc, start, sign, fresh):
    if fresh:
        mc.set_fresh_mode(1)
    st = {"last": 0.0, "n": 0}
    t_cmd = now()

    def tick(t, a):
        if t - st["last"] >= STREAM_INTERVAL_SEC:
            target = list(start)
            target[0] = a[0] + sign * STREAM_AHEAD_DEG
            mc.send_angles(target, SPEED)
            st["last"] = t
            st["n"] += 1
    try:
        samples, guard = run_sampling(mc, start, STREAM_SEC, sign, on_tick=tick)
        t_stop = now()
        after, _ = run_sampling(mc, start, AFTER_STOP_SEC, sign)
    finally:
        if fresh:
            mc.set_fresh_mode(0)
    s = summarize(samples + after, t_cmd, t_stop, sign)
    s["n_sends"] = st["n"]
    return s, samples + after, t_cmd, t_stop


def test_jog(mc, start, sign, speed):
    direction = 1 if sign > 0 else 0
    t_cmd = now()
    try:
        mc.jog_angle(1, direction, speed)
        samples, guard = run_sampling(mc, start, JOG_SEC, sign)
    finally:
        mc.jog_stop()
    t_stop = now()
    after, _ = run_sampling(mc, start, AFTER_STOP_SEC, sign)
    s = summarize(samples + after, t_cmd, t_stop, sign)
    moved = abs(samples[-1][1] - samples[0][1]) if len(samples) >= 2 else 0.0
    s["jog_moved"] = moved > 1.0
    return s, samples + after, t_cmd, t_stop


def show(name, s):
    keys = ["travel_deg", "mean_speed", "median_speed", "max_speed", "speed_cv", "stall_ratio",
            "overshoot_deg", "stop_latency_sec", "arrive_sec", "n_sends", "jog_moved"]
    print("  [{}] ".format(name) + "  ".join("{}={}".format(k, s[k]) for k in keys if k in s), flush=True)
    if "bin_speeds" in s:
        print("        0.1 秒ごとの速さ (°/秒): " + " ".join(str(v) for v in s["bin_speeds"]), flush=True)


def main(argv=None):
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", required=True)
    args = ap.parse_args(argv)

    def on_hup(signum, frame):
        raise Quit()
    signal.signal(signal.SIGHUP, on_hup)

    import pymycobot
    from pymycobot.mycobot import MyCobot
    print("pymycobot {} ({})".format(getattr(pymycobot, "__version__", "?"), pymycobot.__file__), flush=True)
    mc = MyCobot(SERIAL_PORT, BAUDRATE)
    time.sleep(0.5)
    result = {"started_at": time.strftime("%Y-%m-%d %H:%M:%S"), "pymycobot": getattr(pymycobot, "__version__", "?"),
              "params": {"speed": SPEED, "stream_interval": STREAM_INTERVAL_SEC, "stream_ahead": STREAM_AHEAD_DEG,
                         "jog_speeds": list(JOG_SPEEDS), "guard_deg": JOG_GUARD_DEG}, "tests": {}}
    try:
        start = read_angles(mc)
        print("始めの角度: {}".format(start), flush=True)
        if start is None:
            print("角度が読めないので終わります", flush=True)
            return 1
        if max(abs(x - y) for x, y in zip(start, HOME)) > HOME_TOL_DEG:
            ask("HOME から離れています。HOME [0,0,0,0,0,-30] へ送ります")
            mc.send_angles(HOME, SPEED)
            start, ok = wait_arrive(mc, HOME, 15.0)
            if not ok:
                print("HOME に着かないので終わります", flush=True)
                return 1
            time.sleep(0.5)
            start = read_angles(mc)
        result["start"] = start
        print("根元 j0 だけを、始めの角度から 25° 以内で動かします。毎回始めの角度へ戻します。", flush=True)

        print("\n0. 読み取りの速さ (動かさない)", flush=True)
        result["tests"]["read"] = test_read(mc)

        plan = [("A_step_10deg", lambda: test_step(mc, start)),
                ("B_stream", lambda: test_stream(mc, start, +1, False)),
                ("C_stream_fresh", lambda: test_stream(mc, start, -1, True)),
                ("D_jog_speed{}".format(JOG_SPEEDS[0]), lambda: test_jog(mc, start, +1, JOG_SPEEDS[0])),
                ("D_jog_speed{}".format(JOG_SPEEDS[1]), lambda: test_jog(mc, start, -1, JOG_SPEEDS[1]))]
        for name, fn in plan:
            ask("\n{} を始めます".format(name))
            s, samples, t_cmd, t_stop = fn()
            show(name, s)
            result["tests"][name] = dict(s, samples=[[round(t - t_cmd, 4), round(q, 2)] for t, q in samples],
                                         t_stop=None if t_stop is None else round(t_stop - t_cmd, 4))
            back_to_start(mc, start)
    except Quit:
        print("\n止めました (このあとは動かす命令を送りません)", flush=True)
        result["quit"] = True
    finally:
        for f in (getattr(mc, "jog_stop", None),):
            try:
                if f is not None:
                    f()
            except Exception:
                pass
        try:
            mc.set_fresh_mode(0)
        except Exception:
            pass
        with open(args.out, "w") as fp:
            json.dump(result, fp, ensure_ascii=False)
        print("結果: " + args.out, flush=True)
    return 0


if __name__ == "__main__":
    sys.exit(main())