#!/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・今の角度へ send_angles を送ってそのテストを止める。

--set stop (2026-11-05 追加): 1 回目 (11/5 15:34) で, jog_stop のあとも j0 が同じ速さで動き続けたため,
止め方を比べる。jog (speed 5・10) と, 遠い目標への 1 回の send_angles (speed 10・20) を 1 秒動かし,
jog_stop / stop / hold (今の角度へ send_angles) のどれで止まるかと, 止める命令のあとに進んだ角度を測る。
jog_angle の向きは, 1 回目の実機で 1 を送ると j0 が減った (説明書きと逆) ので, JOG_DIR で合わせる。

--set hold (2026-11-05 追加): 2 回目 (15:39) で, jog は send_angles では止まらず (hold が効かない),
stop は効いたり効かなかったり, jog_stop も 1 回目は効かなかった。遠い目標への send_angles は
1 秒で目標に着いてしまい, 途中で止める試しになっていなかった。そこで,
G: 遠い目標 (20°) へ送り, 動いている途中 (0.4〜0.5 秒) で止める (hold / stop)
H: jog speed 10 を jog_stop で止めるのを 3 回 (止まる確かさ) を測る。

--set small (2026-11-05 追加): ゲームパッドの 1 刻み (1°・2°) が動いたり動かなかったりする. どの送り方なら
小さな動きが通るかを測る. 腕を少し前に倒した姿勢 SMALL_POSE (j1 が重さで指令より下がる. 11/5 の実機で 3〜6°) で,
j0 を ±1° と ±2° ずつ, 次の 3 通りで 6 回ずつ送り, 1 秒のうちに 0.5° 以上動いたかを数える.
  S1 send_angles: ほかの関節は「いま読んだ角度」
  S2 send_angles: ほかの関節は「姿勢の指令の角度」(重さで下がった分とずれている. ゲームパッドの今の送り方)
  S3 send_angle: j0 だけ (ほかの関節には指令を送らない)
最後に HOME へ戻す。
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
JOG_DIR = {+1: 0, -1: 1}   # 11/5 の実機: direction=1 で j0 が減った, 0 で増えた
FAR_DEG = 20.0
AFTER_STOP_SEC_STOP_SET = 1.5
FAILS = {"n": 0}
SMALL_POSE = [0.0, 40.0, -20.0, 0.0, -30.0, -30.0]   # 11/5 16:13 のゲームパッドで実際にいた姿勢に近い
SMALL_STEPS = (1.0, 2.0)
SMALL_REPEAT = 6
SMALL_WAIT_SEC = 1.0
SMALL_MOVED_DEG = 0.5
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 not (a and len(a) == 6):
            FAILS["n"] += 1
        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)
                guard_halt(mc, [float(x) for x in a])
                return samples, True
            if on_tick is not None:
                on_tick(t, [float(x) for x in a])
    return samples, False


def guard_halt(mc, a):
    """守りに当たったとき: jog_stop・stop・今の角度への send_angles を全部送る"""
    for f in (lambda: mc.jog_stop(), lambda: mc.stop(), lambda: mc.send_angles(a, SPEED)):
        try:
            f()
        except Exception:
            pass


def halt(mc, how):
    """止め方を 1 つだけ使う. hold は今の角度を読んでそこへ send_angles"""
    if how == "jog_stop":
        mc.jog_stop()
    elif how == "stop":
        mc.stop()
    elif how == "hold":
        a = read_angles(mc, tries=3)
        if a is not None:
            mc.send_angles(a, SPEED)
    else:
        raise ValueError(how)


def stop_set_result(mc, start, samples, after, t_cmd, t_stop, sign):
    s = summarize(samples + after, t_cmd, t_stop, sign)
    s["read_fails"] = FAILS["n"]
    q_end = after[-1][1] if after else None
    s["end_from_start_deg"] = None if q_end is None else round(q_end - start[0], 2)
    # 止める命令から 1.5 秒のあいだに, 最後に 0.1° より動いた時刻 (止まるまでの時間)
    return s


def test_jog_stopset(mc, start, sign, speed, how):
    FAILS["n"] = 0
    t_cmd = now()
    mc.jog_angle(1, JOG_DIR[sign], speed)
    samples, guard = run_sampling(mc, start, JOG_SEC, sign)
    t_stop = now()
    if not guard:
        halt(mc, how)
    after, _ = run_sampling(mc, start, AFTER_STOP_SEC_STOP_SET, sign)
    s = stop_set_result(mc, start, samples, after, t_cmd, t_stop, sign)
    s.update(how=how, speed=speed, guard=guard)
    return s, samples + after, t_cmd, t_stop


def test_far_stopset(mc, start, sign, speed, how, sec=JOG_SEC):
    FAILS["n"] = 0
    target = list(start)
    target[0] = start[0] + sign * FAR_DEG
    t_cmd = now()
    mc.send_angles(target, speed)
    samples, guard = run_sampling(mc, start, sec, sign)
    t_stop = now()
    if not guard:
        halt(mc, how)
    after, _ = run_sampling(mc, start, AFTER_STOP_SEC_STOP_SET, sign)
    s = stop_set_result(mc, start, samples, after, t_cmd, t_stop, sign)
    s.update(how=how, speed=speed, guard=guard)
    return s, samples + after, t_cmd, t_stop


def test_small(mc, how, step):
    """j0 を ±step ずつ SMALL_REPEAT 回. 毎回 送る前に読み, 1 秒のあいだ読んで 0.5° 以上動いたかを見る."""
    trials = []
    sign = +1
    for i in range(SMALL_REPEAT):
        stdin_quit()
        q = read_angles(mc)
        if q is None:
            trials.append({"i": i, "error": "read_failed"})
            continue
        want = q[0] + sign * step
        if how == "S1":
            target = list(q)
            target[0] = want
            mc.send_angles(target, SPEED)
        elif how == "S2":
            target = list(SMALL_POSE)
            target[0] = want
            mc.send_angles(target, SPEED)
        else:
            mc.send_angle(1, want, SPEED)
        t0 = now()
        moved_at, q_end = None, q[0]
        while now() - t0 < SMALL_WAIT_SEC:
            a = read_angles(mc, tries=3)
            if a is not None:
                q_end = a[0]
                if moved_at is None and abs(a[0] - q[0]) >= SMALL_MOVED_DEG:
                    moved_at = now() - t0
        moved = (q_end - q[0]) * sign
        trials.append({"i": i, "sign": sign, "q0": round(q[0], 2), "want": round(want, 2),
                       "moved_deg": round(moved, 2), "moved": moved_at is not None,
                       "moved_at_sec": None if moved_at is None else round(moved_at, 2),
                       "others_gap_deg": round(max(abs(q[k] - SMALL_POSE[k]) for k in range(1, 5)), 2)})
        sign = -sign
    n_ok = sum(1 for t in trials if t.get("moved"))
    print("  [{} {}°] 動いた {}/{}  動いた量: {}".format(how, step, n_ok, len(trials),
          " ".join(str(t.get("moved_deg")) for t in trials)), flush=True)
    return {"how": how, "step": step, "n_moved": n_ok, "n": len(trials), "trials": trials}


def run_small_set(mc, result):
    ask("\n腕を少し前に倒した姿勢 {} へ動かします".format(SMALL_POSE))
    mc.send_angles(SMALL_POSE, SPEED)
    a, ok = wait_arrive(mc, SMALL_POSE, 15.0)
    time.sleep(1.0)
    a = read_angles(mc)
    print("  着いた角度: {}  (指令との差 j1 {:+.2f}°)".format(a, None if a is None else a[1] - SMALL_POSE[1]), flush=True)
    result["small_pose_measured"] = a
    out = {}
    for how in ("S1", "S2", "S3"):
        for step in SMALL_STEPS:
            ask("\n{} {}° を始めます".format(how, step))
            out["{}_{}".format(how, int(step))] = test_small(mc, how, step)
    result["tests"]["small"] = out
    ask("\nHOME へ戻します")
    mc.send_angles(HOME, SPEED)
    wait_arrive(mc, HOME, 15.0)


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",
            "read_fails", "end_from_start_deg", "guard"]
    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)
    ap.add_argument("--set", choices=["base", "stop", "hold", "small"], default="base")
    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,
                         "set": args.set, "jog_dir": {"+1": JOG_DIR[+1], "-1": JOG_DIR[-1]}}, "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_stop = [
            ("E_jog5_jogstop", lambda: test_jog_stopset(mc, start, +1, 5, "jog_stop")),
            ("E_jog5_stop", lambda: test_jog_stopset(mc, start, -1, 5, "stop")),
            ("E_jog5_hold", lambda: test_jog_stopset(mc, start, +1, 5, "hold")),
            ("E_jog10_stop", lambda: test_jog_stopset(mc, start, -1, 10, "stop")),
            ("E_jog10_hold", lambda: test_jog_stopset(mc, start, +1, 10, "hold")),
            ("F_far10_stop", lambda: test_far_stopset(mc, start, -1, 10, "stop")),
            ("F_far10_hold", lambda: test_far_stopset(mc, start, +1, 10, "hold")),
            ("F_far20_hold", lambda: test_far_stopset(mc, start, -1, 20, "hold"))]
        plan_hold = [
            ("G_far5_hold_0.5s", lambda: test_far_stopset(mc, start, +1, 5, "hold", 0.5)),
            ("G_far10_hold_0.5s", lambda: test_far_stopset(mc, start, -1, 10, "hold", 0.5)),
            ("G_far20_hold_0.4s", lambda: test_far_stopset(mc, start, +1, 20, "hold", 0.4)),
            ("G_far10_stop_0.5s", lambda: test_far_stopset(mc, start, -1, 10, "stop", 0.5)),
            ("H_jog10_jogstop_1", lambda: test_jog_stopset(mc, start, +1, 10, "jog_stop")),
            ("H_jog10_jogstop_2", lambda: test_jog_stopset(mc, start, -1, 10, "jog_stop")),
            ("H_jog10_jogstop_3", lambda: test_jog_stopset(mc, start, +1, 10, "jog_stop"))]
        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]))]
        if args.set == "small":
            run_small_set(mc, result)
            plan = []
        elif args.set == "stop":
            plan = plan_stop
        elif args.set == "hold":
            plan = plan_hold
        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())
