# -*- coding: utf-8 -*-
"""Compute each axis's real free travel by sweeping the assembled solids.

For joint Jk the moving set is every part whose kinematic group is Jk or a descendant
of it, rotated about the joint axis; the static set is everything else.  A pair can only
jam if the analytically rotated bounding boxes overlap, so the sweep is
   box prune -> boolean intersection volume -> jam
and the first angle whose overlap exceeds VOL_MIN is the hard stop.  Where a whole turn
stays clear the axis is reported continuous, which is what the viewer's slider needs.

The jaw slides (P10/P11) are excluded: they are constrained along their rail, not by the
tree, and their own crash limit is the crank travel.

Run:  python travel_sweep.py [J1 J3 ...]      -> docs/travel.json
"""
import json
import math
import os
import sys
import time

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
try:
    sys.stdout.reconfigure(encoding="utf-8")
except Exception:
    pass
import cadquery as cq
from armcore import WEB, DOCS, OUT, place
from arm_build import R, bb, common_vol
from OCP.gp import gp_Pnt

VOL_MIN = 2.0            # mm3 of interpenetration that counts as a jam
COARSE = 3.0             # deg, scan step
FINE = 0.25              # deg, bisect resolution
TREE = ["root", "J1", "J2", "J3", "J4", "J5", "J6"]


def load():
    meta = json.load(open(WEB + "/mesh_arm.json", encoding="utf-8"))
    joints = {j["id"]: j for j in meta["joints"]}
    parts = {}
    for p in meta["parts"]:
        if p["jaw"]:
            continue
        f = "%s/parts_step/%s.step" % (OUT, p["name"])
        if not os.path.exists(f):
            print("  missing %s" % f)
            continue
        parts[p["name"]] = (cq.importers.importStep(f).val(), p["grp"])
    return joints, parts


def corners(lo, hi):
    return [(x, y, z) for x in (lo[0], hi[0]) for y in (lo[1], hi[1]) for z in (lo[2], hi[2])]


def rot_box(c8, ang, axis, pt):
    t = R(ang, axis[0], axis[1], axis[2], pt[0], pt[1], pt[2])
    ps = []
    for c in c8:
        q = gp_Pnt(*c).Transformed(t)
        ps.append((q.X(), q.Y(), q.Z()))
    return ([min(p[i] for p in ps) for i in range(3)], [max(p[i] for p in ps) for i in range(3)])


def overlap(mov, sta, ang, axis, pt):
    """Interpenetration volume (mm3) with the moving set rotated by ang about the axis.

    The rotation must reach the GEOMETRY, not just the prune box - transforming only the
    boxes leaves every boolean at the assembled pose, which reports any axis as clear."""
    t = R(ang, axis[0], axis[1], axis[2], pt[0], pt[1], pt[2])
    boxes = {n: rot_box(m["c8"], ang, axis, pt) for n, m in mov.items()}
    cands = []
    for sn, s in sta.items():
        for mn, b1 in boxes.items():
            b2 = s["box"]
            if any(b1[1][k] < b2[0][k] + 0.05 or b2[1][k] < b1[0][k] - 0.05 for k in range(3)):
                continue
            cands.append((mn, sn))
    if not cands:
        return 0.0
    rot = {}
    for mn in {m for m, _ in cands}:
        rot[mn] = place(mov[mn]["shape"], t)
    tot = 0.0
    for mn, sn in cands:
        v = common_vol(rot[mn], sta[sn]["shape"])
        if v > VOL_MIN:
            tot += v
    return tot


def refine(mov, sta, axis, pt, a_free, a_jam):
    while abs(a_jam - a_free) > FINE:
        mid = 0.5 * (a_free + a_jam)
        if overlap(mov, sta, mid, axis, pt) > VOL_MIN:
            a_jam = mid
        else:
            a_free = mid
    return a_free


def sweep_one(jid, joints, parts):
    j = joints[jid]
    k = TREE.index(jid)
    down = set(TREE[k:])
    axis, pt = j["d"], j["rear"]
    mov = {n: dict(shape=s, box=bb(s), c8=corners(*bb(s)))
           for n, (s, g) in parts.items() if g in down}
    sta = {n: dict(shape=s, box=bb(s)) for n, (s, g) in parts.items() if g not in down}
    zero = overlap(mov, sta, 0.0, axis, pt)
    full = True
    a = 0.0
    while a < 360.0 - 1e-6:
        if overlap(mov, sta, a, axis, pt) > VOL_MIN:
            full = False
            break
        a += COARSE
    if full:
        return dict(axis=jid, lo=-180.0, hi=180.0, continuous=True, zero=round(zero, 1),
                    note="整周无干涉，可 360° 连续旋转")
    pos = _limit(mov, sta, axis, pt, +1)
    neg = _limit(mov, sta, axis, pt, -1)
    hi = min(pos, 180.0)
    lo = max(neg, -180.0)
    return dict(axis=jid, lo=round(lo, 1), hi=round(hi, 1), continuous=False,
                zero=round(zero, 1),
                note="正向 +%.0f° / 反向 %.0f°，总行程 %.0f°" % (hi, lo, hi - lo))


def _limit(mov, sta, axis, pt, sign):
    """How far the joint turns in `sign` direction before something interpenetrates."""
    step = COARSE
    a_free = 0.0
    while step < 360.0:
        ang = sign * step
        if overlap(mov, sta, ang % 360.0, axis, pt) > VOL_MIN:
            return refine(mov, sta, axis, pt, a_free, ang)
        a_free = ang
        step += COARSE
    return 180.0 * sign


def main():
    t0 = time.time()
    joints, parts = load()
    print("loaded %d rigid parts, groups=%s" % (len(parts), sorted({g for _, g in parts.values()})))
    want = sys.argv[1:] or [k for k in ("J1", "J2", "J3", "J4", "J5", "J6") if k in joints]
    rows = []
    for jid in want:
        r = sweep_one(jid, joints, parts)
        rows.append(r)
        print("%-3s lo=%7.1f hi=%7.1f continuous=%-5s zero=%6.1f mm3  %-34s [%.0fs]"
              % (jid, r["lo"], r["hi"], r["continuous"], r["zero"], r["note"],
                 time.time() - t0), flush=True)
    json.dump(rows, open(DOCS + "/travel.json", "w", encoding="utf-8"),
              ensure_ascii=False, indent=1)
    print("wrote docs/travel.json  (%.0fs)" % (time.time() - t0))


if __name__ == "__main__":
    main()