"""Door robot arm: per-joint sweeps with the real meshes (stage geo_sweep -> build/geo_sweep.json).

Each joint is turned through the range the door animation uses, per version (A, B: their own poses since stage 2),
plus a margin of MARGIN degrees past both ends, and always through the rest pose (0). The ranges come from the
closed-form inverse kinematics over the doors' own animation (dra_ik.py, the same maths as the rig's drivers). All
other joints stay at rest (the ref1 pose), on the showroom arms. Everything downstream of the joint moves as one rigid
group and is tested against the parts upstream of it. J1 turns the whole circle.
- J3 also moves the J3 cylinder and its piston the way the rig does (cylinder turned to the knob, piston slid).
- J5: the wrist (J5) and the flange roll (J6) move together on the door path. At the deep wrist bends (frames 9-15
  and 100-112) the path holds J6 at 45-52 degrees, so no finger is in the bend plane there; the wrist block is shaped
  for that (dra_parts_wrist.py). The J5 sweep therefore turns J5 past the path's ends with J6 as the path has it
  there, and turns J6 10 degrees either side of the path at the path's points (one margin at a time, like the other
  joints). Information only: 'J5 limit with J6 at rest' (how far J5 bends with J6 at 0) and 'J5 margin at the
  deepest bend' (with J6 on the path and 10 degrees off it).
- The grip closes the fingers by the design's angles onto the knob fixture on the tool axis.
This is not the rig's check (rig_check runs the real path, all joints together, every frame); it finds geometry
problems early and shows the margins.
"""
import math
import json
import os
import bpy
from mathutils import Vector, Matrix
from mathutils.bvhtree import BVHTree
import dra_audit as au
import dra_rigdata as rd
import dra_ik

MARGIN = 3.0
STEP = 3.0
Q6_BAND = 10.0
SWEEPS = {
    # joint: (moving link, static links, axis in A, origin in A)
    "J1": (None, ["carriage", "carriage pinion"], 'Z', Vector((0, 0, 0))),
    "J2": ("lower arm", ["turret", "carriage", "carriage pinion"], 'Y', rd.J2),
    "J3": ("upper arm", ["lower arm", "turret"], 'Y', rd.J3),
    "J4": ("forearm", ["upper arm", "lower arm"], 'X', rd.J4),
    "J5": ("wrist head", ["forearm", "upper arm"], 'Y', rd.WC),
    "J6": ("flange", ["wrist head", "forearm"], 'X', rd.WC),
}
# channel sign: the angle in degrees is applied as rotation_euler about the axis with these signs (design.json):
# J2 rot y = -q2, J3 rot y = -q3, J5 rot y = +q5, J4/J6 rot x = +q, J1 rot z = +q1
SIGN = {"J1": 1, "J2": -1, "J3": -1, "J4": 1, "J5": 1, "J6": 1}


def downstream(name):
    out = [name]
    changed = True
    while changed:
        changed = False
        for s in rd.LINKS:
            if s["parent"] in out and s["name"] not in out:
                out.append(s["name"])
                changed = True
    return out


def world_geo(ob, dg):
    ev = ob.evaluated_get(dg)
    M = ev.matrix_world
    me = ob.data
    me.calc_loop_triangles()
    lay = me.attributes.get('part')
    ids = list(me.get("parts", []))
    dec = list(me.get("decals", []))
    verts = [M @ v.co for v in me.vertices]
    tris, parts = [], []
    for lt in me.loop_triangles:
        pv = lay.data[lt.polygon_index].value if lay is not None else -1
        if dec and 0 <= pv < len(dec) and dec[pv]:
            continue
        tris.append(tuple(lt.vertices))
        parts.append(ids[pv] if 0 <= pv < len(ids) else "?")
    return verts, tris, parts


def test(movers, statics, M, extra=None):
    """movers / statics: {name: (verts, tris, parts)}; M: world matrix applied to the movers; extra: per-mover matrix"""
    hits = []
    st = {n: (g[0], g[1], BVHTree.FromPolygons(g[0], g[1], all_triangles=True), g[2]) for n, g in statics.items() if g[1]}
    for mn, (v, t, p) in movers.items():
        Mm = extra[mn] if extra and mn in extra else M
        vv = [Mm @ x for x in v]
        tr = BVHTree.FromPolygons(vv, t, all_triangles=True)
        for sn, (sv, stt, sbvh, sp) in st.items():
            pl = []
            n = au.crossing((vv, t, tr), (sv, stt, sbvh), pl)
            if n:
                det = sorted(set("%s x %s" % (p[i], sp[j]) for i, j in pl))
                hits.append([mn, sn, n, det[:8]])
    return hits


def rot_about(axis, origin_w, deg):
    ax = {'X': Vector((1, 0, 0)), 'Y': Vector((0, 1, 0)), 'Z': Vector((0, 0, 1))}[axis]
    return Matrix.Translation(origin_w) @ Matrix.Rotation(math.radians(deg), 4, ax) @ Matrix.Translation(-origin_w)


def samples(lo, hi, step=STEP):
    """lo..hi (degrees) in steps, both ends exact, rest (0) included"""
    out = {round(lo, 2), round(hi, 2)}
    a = math.ceil(lo / step) * step
    while a < hi:
        out.add(round(a, 2))
        a += step
    if lo <= 0 <= hi:
        out.add(0.0)
    return sorted(out)


def design():
    with open(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "design.json")) as f:
        return json.load(f)


def door_paths():
    sf = bpy.data.scenes["DoorArm | Ship fit"]
    wm = bpy.context.window_manager
    if wm.windows:
        wm.windows[0].scene = sf
    out = {}
    for key, nm in (("AFT", "Moonpool_Door_Aft"), ("FORE", "Moonpool_Door_Fore")):
        out[key] = dra_ik.door_path(sf, bpy.data.objects[nm])
    sf.frame_set(1)
    return out


def run():
    D = design()
    dps = door_paths()
    sc = bpy.data.scenes["DoorArm | Showroom"]
    wm = bpy.context.window_manager
    if wm.windows:
        wm.windows[0].scene = sc
    dg = sc.view_layers[0].depsgraph
    dg.update()
    rep = {"margin_deg": MARGIN, "step_deg": STEP, "q6_band_deg": Q6_BAND, "ranges_used": {}, "ranges_swept": {}, "info": {}}
    for v in ("A", "B"):
        p = dra_ik.path(D, v, dps)
        used = dra_ik.ranges_deg(p)
        rep["ranges_used"][v] = {j: [round(a, 2), round(b, 2)] for j, (a, b) in used.items()}
        prefix = "Showroom | arm %s | " % v
        root = bpy.data.objects["Showroom | arm " + v]
        R0 = root.matrix_world.copy()
        objs = {o["link"]: o for o in sc.objects if o.name.startswith(prefix) and o.type == 'MESH' and "link" in o}
        geo = {n: world_geo(o, dg) for n, o in objs.items()}
        swept = {}
        for j, (mv, stat, axis, org) in SWEEPS.items():
            if j == "J1":
                lo, hi = -180.0, 180.0
                angles = samples(lo, hi, 15.0)
            else:
                lo = min(used[j][0] - MARGIN, 0.0)
                hi = max(used[j][1] + MARGIN, 0.0)
                angles = samples(lo, hi)
            swept[j] = [round(lo, 2), round(hi, 2)]
            names = downstream(mv) if mv else [n for n in objs if n not in ("carriage", "carriage pinion")]
            movers = {n: geo[n] for n in names if n in geo}
            statics = {n: geo[n] for n in stat if n in geo}
            ow = R0 @ org
            res = []
            if j == "J5":
                res = _j5_coupled(v, p, geo, names, statics, R0, swept, rep["info"])
                rep["info"]["J5 limit with J6 at rest, arm %s" % v] = _j5_free(geo, names, statics, R0, lo)
                rep["J5 arm %s" % v] = res
                continue
            for a in angles:
                M = rot_about(axis, ow, a * SIGN[j])
                extra = None
                if j == "J3":
                    extra = _j3_cylinder(R0, a)
                    for n in ("J3 cylinder", "J3 piston"):
                        if n in geo:
                            movers[n] = geo[n]
                h = test(movers, statics, M, extra)
                if j == "J3":                       # the cylinder and piston also against the upper arm group
                    up = {n: geo[n] for n in names if n in geo}
                    cp = {n: geo[n] for n in ("J3 cylinder", "J3 piston") if n in geo}
                    moved_up = {n: ([M @ x for x in g[0]], g[1], g[2]) for n, g in up.items()}
                    h += test(cp, moved_up, Matrix.Identity(4), extra)
                for x in h:
                    res.append([a] + x)
            rep["%s arm %s" % (j, v)] = res
        rep["ranges_swept"][v] = swept
        # grip (version A): close the fingers by g = 0..1
        if v == "A":
            res = []
            hub = {n: geo[n] for n in ("tool A hub",) if n in geo}
            fx = _fixture_on_tool(R0)
            if fx:
                hub["door fixture A (on the tool axis)"] = fx
            for g in (0.0, 0.25, 0.5, 0.75, 1.0):
                fm = {}
                for w in ("top", "left", "bottom", "right"):
                    k, pp = rd.finger_points(w)
                    kw, pw = R0 @ k, R0 @ pp
                    ax, sg = {"top": ('Y', 1), "bottom": ('Y', -1), "left": ('Z', -1), "right": ('Z', 1)}[w]
                    Mp = rot_about(ax, kw, sg * 14.7 * g)
                    pw2 = Mp @ pw
                    Md = rot_about(ax, pw2, sg * 41.0 * g) @ Mp
                    fm["finger %s proximal" % w] = Mp
                    fm["finger %s distal" % w] = Md
                fingers = {n: geo[n] for n in fm if n in geo}
                h = test(fingers, hub, Matrix.Identity(4), fm)
                names_f = list(fingers)
                for i, a_ in enumerate(names_f):
                    for b_ in names_f[i + 1:]:
                        moved_b = {b_: ([fm[b_] @ x for x in geo[b_][0]], geo[b_][1], geo[b_][2])}
                        h += test({a_: geo[a_]}, moved_b, fm[a_])
                for x in h:
                    res.append([g] + x)
            rep["grip arm A"] = res
    rep["summary"] = {k: len(val) for k, val in rep.items() if k.startswith(("J", "grip"))}
    return rep


def _wrist_mats(R0, q5, q6, names):
    """world matrices of the wrist group: the head turns by q5 about J5 (Y at WC), the flange and everything on it
    also by q6 about the tool axis (X at WC, inside the head)"""
    Rh = rot_about('Y', R0 @ rd.WC, q5 * SIGN["J5"])
    Rf = Rh @ rot_about('X', R0 @ rd.WC, q6 * SIGN["J6"])
    return {n: (Rh if n == "wrist head" else Rf) for n in names}


def _j5_coupled(v, p, geo, names, statics, R0, swept, info):
    """the wrist on the door path, one margin at a time (as the other joints' sweeps): J5 turned MARGIN past the
    path's deepest and shallowest bends with J6 as the path has it there, and J6 turned Q6_BAND either side of the
    path at every other path point. Both margins at once (the corner) is measured as information."""
    movers = {n: geo[n] for n in names if n in geo}
    pts = sorted({(round(math.degrees(q["J5"]), 1), round(math.degrees(q["J6"]), 1)) for _, _, q in p})
    q5s = [a for a, _ in pts]
    deep = [pt for pt in pts if pt[0] <= min(q5s) + 0.5]
    shallow = [pt for pt in pts if pt[0] >= max(q5s) - 0.5]
    tests = set()
    for (a, b) in pts[::2] + deep + shallow:
        for db in (-Q6_BAND, 0.0, Q6_BAND):
            tests.add((a, round(b + db, 1)))
    for k in (1, 2, 3):
        e = MARGIN * k / 3
        for (a, b) in deep:
            tests.add((round(a - e, 1), b))
        for (a, b) in shallow:
            tests.add((round(a + e, 1), b))
    swept["J5"] = [round(min(q5s) - MARGIN, 2), round(max(q5s) + MARGIN, 2)]
    swept["J5_with_J6"] = "J6 as on the door path; J6 also +-%.0f deg along the path" % Q6_BAND
    res = []
    for (a, b) in sorted(tests):
        h = test(movers, statics, Matrix.Identity(4), _wrist_mats(R0, a, b, list(movers)))
        for x in h:
            res.append([a, b] + x)
    # information: at the deepest bend, how much further J5 goes with J6 on the path and with J6 Q6_BAND off it
    corner = {}
    for (a, b) in deep[:2]:
        for db in (-Q6_BAND, 0.0, Q6_BAND):
            x = a
            while x > a - 15.0:
                if test(movers, statics, Matrix.Identity(4), _wrist_mats(R0, x - 0.5, b + db, list(movers))):
                    break
                x -= 0.5
            corner["J6 %+.1f" % (b + db)] = dict(path_deepest_J5=a, free_to_J5=round(x, 1), margin_deg=round(a - x, 1))
    info["J5 margin at the deepest bend, arm %s" % v] = corner
    return res


def _j5_free(geo, names, statics, R0, lo):
    """how far J5 bends (toward the forearm's +Z side, q5 < 0) with J6 at rest before anything touches"""
    movers = {n: geo[n] for n in names if n in geo}
    a = 0.0
    last_ok = 0.0
    while a >= -120.0:
        h = test(movers, statics, Matrix.Identity(4), _wrist_mats(R0, a, 0.0, list(movers)))
        if h:
            return dict(free_to_deg=last_ok, first_contact_deg=a, parts=sorted(set(d for x in h for d in x[3]))[:6])
        last_ok = a
        a -= 0.5
    return dict(free_to_deg=last_ok)


def _j3_cylinder(R0, q3):
    """world matrices of the cylinder and piston for elbow angle q3 (lower arm at rest)"""
    T = rd.TRUNNION
    K0 = rd.KNOB
    J3 = rd.J3
    # the knob turns with the upper arm: rotation_euler.y = -q3 about J3
    Rk = Matrix.Translation(J3) @ Matrix.Rotation(math.radians(-q3), 4, 'Y') @ Matrix.Translation(-J3)
    K = Rk @ K0
    a0 = math.atan2(K0.z - T.z, K0.x - T.x)
    a1 = math.atan2(K.z - T.z, K.x - T.x)
    d = a1 - a0                                   # angle change in the x-z angle sense
    Rc = Matrix.Translation(T) @ Matrix.Rotation(-d, 4, 'Y') @ Matrix.Translation(-T)
    L0 = (K0 - T).length
    L = (K - T).length
    u = (K - T).normalized()
    Mp = Matrix.Translation(u * (L - L0)) @ Rc
    return {"J3 cylinder": R0 @ Rc @ R0.inverted(), "J3 piston": R0 @ Mp @ R0.inverted()}


def _fixture_on_tool(R0):
    """the AFT door's fixture A moved into the showroom arm's tool frame at rest: its face normal along -x', the door
    face 0.63 m ahead of WC (design.json tools.A)"""
    ob = bpy.data.objects.get("DoorArm_AFT | door fixture A")
    if ob is None:
        return None
    d = design()
    n = Vector(d["doors"]["AFT"]["normal_local"]).normalized()
    xa = Vector((1, 0, 0))
    ya = n.cross(xa).normalized()
    me = ob.data
    me.calc_loop_triangles()
    face = rd.WC + Vector((dra_ik.D_T["A"], 0, 0))
    verts = []
    for v in me.vertices:
        u, w2, w = v.co.dot(xa), v.co.dot(ya), v.co.dot(n)
        a = face + Vector((0, 1, 0)) * u + Vector((0, 0, -1)) * w2 + Vector((-1, 0, 0)) * w
        verts.append(R0 @ a)
    tris = [tuple(lt.vertices) for lt in me.loop_triangles]
    return verts, tris, ["fixture A"] * len(tris)