"""Check the robot cell: drivers, every moving part through its whole travel, the robot's work cycle, topology.
Nothing is saved.

    blender.exe -b ../RobotCell.blend --python rc_check.py      -> build/check_report.json
    blender.exe -b ../RobotCell_DoorArm.blend --python rc_check.py -- arm=RC_DoorArm_Rig mesh=SK_RobotCell_DoorArm tag=doorarm
                                                                -> build/check_report_doorarm.json
A rig that carries its own "axes", "limits_deg" and "skip_pairs" (the door-arm copy) is checked by those; a J3
cylinder (bones J3_Cylinder, J3_Piston, J3_Pin) also gets its stroke checked against its barrel.

* the two leaves (each alone and both together), the side door, the turntable plate (a full turn) and the emergency
  stop through their travel in steps against the cell's body and each other (triangle overlap, BVH);
* the robot (posed mesh, world space) at every second frame of its work cycle against the body and the moving parts,
  and against itself: its rigid parts per bone, pairwise (the hoses are soft and left out; A6/Tool and Tool/Tool_Spin
  are bolted or bearing contacts and left out);
* each joint swept alone over its limits from the picture's pose (self only), to show where the limits bite;
* topology audit of every mesh.
"""
import bpy, bmesh, json, os, math, sys
from mathutils import Vector
from mathutils.bvhtree import BVHTree

HERE = os.path.dirname(os.path.abspath(__file__))
exec(open(os.path.join(HERE, "rc_lib.py")).read())
exec(open(os.path.join(HERE, "rc_robot.py")).read())

ARGS = dict(a.split("=", 1) for a in (sys.argv[sys.argv.index("--") + 1:] if "--" in sys.argv else []) if "=" in a)
sc = bpy.context.scene
root = bpy.data.objects["RobotCell_Root"]
arm = bpy.data.objects[ARGS.get("arm", "RC_Robot_Rig")]
rob = bpy.data.objects[ARGS.get("mesh", "SK_RobotCell_Robot")]
TAG = ARGS.get("tag", "")
if "limits_deg" in arm.keys():
    LIMITS = {k: tuple(v) for k, v in json.loads(arm["limits_deg"]).items()}
    AXIS = json.loads(arm["axes"])
    JOINTS = list(LIMITS)

    def set_pose(a, th):
        for nm, ax in AXIS.items():
            e = [0.0, 0.0, 0.0]
            e["XYZ".index(ax)] = th.get(nm, 0.0)
            a.pose.bones[nm].rotation_euler = e
body = bpy.data.objects["SM_RobotCell_Body"]
MOV = ["SM_RobotCell_DoorLeft", "SM_RobotCell_DoorRight", "SM_RobotCell_DoorSide", "SM_RobotCell_TurntablePlate",
       "SM_RobotCell_EStop"]


def bvh_of(ob, faces=None):
    dg = bpy.context.evaluated_depsgraph_get()
    ev = ob.evaluated_get(dg)
    me = ev.to_mesh()
    mw = ev.matrix_world
    co = [mw @ v.co for v in me.vertices]
    polys = [list(p.vertices) for p in me.polygons] if faces is None else [list(me.polygons[i].vertices) for i in faces]
    ev.to_mesh_clear()
    return BVHTree.FromPolygons(co, polys, epsilon=0.0)


def update(frame=None):
    root.update_tag()
    arm.update_tag()
    if frame is not None:
        sc.frame_set(frame)
    else:
        sc.frame_set(sc.frame_current)
    bpy.context.evaluated_depsgraph_get().update()


def first_faces(ob, pairs, k=6, other=None):
    """where: material names and centres of the first faces of a clash"""
    me = ob.data
    names = [m.name for m in me.materials]
    out = []
    for i, _ in pairs[:k]:
        p = me.polygons[i]
        out.append((names[p.material_index], [round(c, 3) for c in ob.matrix_world @ p.center]))
    return out


rep = {"drivers": {}, "travel": {}, "robot_cycle": {}, "robot_sweep": {}, "audit": {}}
# --- drivers
for ob in [bpy.data.objects[n] for n in MOV]:
    if ob.animation_data:
        for fc in ob.animation_data.drivers:
            d = fc.driver
            rep["drivers"]["%s.%s[%d]" % (ob.name, fc.data_path, fc.array_index)] = dict(
                valid=d.is_valid, simple=d.is_simple_expression, expression=d.expression)
# no animation while sliders are tested by hand
saved = {o.name: o.animation_data.action for o in (root, arm) if o.animation_data}
root.animation_data.action = None
for k in ("door_left", "door_right", "side_door", "turntable", "estop"):
    root[k] = 0.0
update(1)
tb = bvh_of(body)


def sweep(name, prop, values, others=()):
    ob = bpy.data.objects[name]
    res = dict(steps=len(values), clashes=[])
    for v in values:
        root[prop] = v
        update()
        t = bvh_of(ob)
        pairs = t.overlap(tb)
        hit = []
        if pairs:
            hit.append(("body", len(pairs), first_faces(body, [(j, i) for i, j in pairs])))
        for on in others:
            p2 = t.overlap(bvh_of(bpy.data.objects[on]))
            if p2:
                hit.append((on, len(p2)))
        if hit:
            res["clashes"].append(dict(value=round(v, 4), hits=hit))
    root[prop] = 0.0
    update()
    res["clash_steps"] = len(res["clashes"])
    res["clashes"] = res["clashes"][:4]
    return res


steps = lambda n, a=0.0, b=1.0: [a + (b - a) * i / n for i in range(n + 1)]
rep["travel"]["door_left"] = sweep("SM_RobotCell_DoorLeft", "door_left", steps(56), ["SM_RobotCell_DoorRight"])
rep["travel"]["door_right"] = sweep("SM_RobotCell_DoorRight", "door_right", steps(56), ["SM_RobotCell_DoorLeft"])
root["door_left"] = 1.0
rep["travel"]["door_right_with_left_open"] = sweep("SM_RobotCell_DoorRight", "door_right", steps(28), ["SM_RobotCell_DoorLeft"])
root["door_left"] = 0.0
rep["travel"]["side_door"] = sweep("SM_RobotCell_DoorSide", "side_door", steps(50))
rep["travel"]["turntable"] = sweep("SM_RobotCell_TurntablePlate", "turntable", steps(36, 0.0, 360.0))
rep["travel"]["estop"] = sweep("SM_RobotCell_EStop", "estop", steps(14))

# --- the robot: rigid parts per bone (faces whose vertices all carry one bone at weight 1, not rubber)
vg = {g.index: g.name for g in rob.vertex_groups}
bone_of = {}
for v in rob.data.vertices:
    ws = [(g.weight, vg[g.group]) for g in v.groups if g.weight > 0.0]
    if len(ws) == 1 and ws[0][0] > 0.999:
        bone_of[v.index] = ws[0][1]
mats = [m.name for m in rob.data.materials]
parts = {}
for p in rob.data.polygons:
    if mats[p.material_index] in ("M_RC_Rubber", "M_DoorArm_Cable"):
        continue
    bs = {bone_of.get(i) for i in p.vertices}
    if len(bs) == 1 and None not in bs:
        parts.setdefault(bs.pop(), []).append(p.index)
SKIP = {frozenset(("A6", "Tool")), frozenset(("Tool", "Tool_Spin"))}
if "skip_pairs" in arm.keys():
    SKIP = {frozenset(p) for p in json.loads(arm["skip_pairs"])}
CYL = "J3_Cylinder" in arm.pose.bones


def stroke():
    """the cylinder's length (trunnion to rod eye) now"""
    pb = arm.pose.bones
    return (pb["J3_Pin"].head - pb["J3_Cylinder"].head).length


if CYL:
    # the barrel along the cylinder's axis and the rod's length, measured from the bind-pose mesh
    bb = arm.data.bones
    T0, K0 = bb["J3_Cylinder"].head_local, bb["J3_Pin"].head_local
    L0 = (K0 - T0).length
    u = (K0 - T0) / L0
    gi = {g.name: g.index for g in rob.vertex_groups}
    tb_, rp_ = [], []
    for v in rob.data.vertices:
        gs = {g.group: g.weight for g in v.groups}
        if gs.get(gi["J3_Cylinder"], 0) > 0.999:
            tb_.append((v.co - T0).dot(u))
        elif gs.get(gi["J3_Piston"], 0) > 0.999:
            rp_.append((K0 - v.co).dot(u))
    BARREL = (min(tb_), max(tb_))
    ROD = max(rp_)
    STROKE = dict(rest_m=round(L0, 4), barrel_m=[round(x, 4) for x in BARREL], rod_m=round(ROD, 4),
                  allowed_m=[round(BARREL[0] + ROD + 0.01, 4), round(BARREL[1] + ROD - 0.02, 4)], seen_m=[9.0, 0.0])

    def note_stroke():
        L = stroke()
        STROKE["seen_m"] = [round(min(STROKE["seen_m"][0], L), 4), round(max(STROKE["seen_m"][1], L), 4)]
names = sorted(parts)


def robot_state(with_cell=True):
    t_all = bvh_of(rob)
    out = []
    if with_cell:
        p = t_all.overlap(tb)
        if p:
            out.append(("body", len(p), first_faces(rob, p, 4)))
        for n in MOV:
            p = t_all.overlap(bvh_of(bpy.data.objects[n]))
            if p:
                out.append((n, len(p)))
    T = {n: bvh_of(rob, parts[n]) for n in names}
    for i, a in enumerate(names):
        for b in names[i + 1:]:
            if frozenset((a, b)) in SKIP:
                continue
            p = T[a].overlap(T[b])
            if p:
                out.append(("self %s/%s" % (a, b), len(p)))
    return out


# the work cycle with its own animation
root.animation_data.action = saved.get(root.name)
arm.animation_data.action = saved.get(arm.name)
cyc = []
for f in range(1, sc.frame_end + 1, 2):
    update(f)
    if CYL:
        note_stroke()
    hit = robot_state()
    if hit:
        cyc.append(dict(frame=f, hits=hit))
rep["robot_cycle"] = dict(frames_tested=len(range(1, sc.frame_end + 1, 2)), clash_frames=len(cyc), first=cyc[:6])
# each joint alone over its limits from the picture's pose (self only)
arm.animation_data.action = None
pic = json.loads(arm["poses_deg"])["picture"]
for j in JOINTS:
    lo, hi = LIMITS[j]
    bad = []
    for a in steps(36, lo, hi):
        th = {k: math.radians(v) for k, v in pic.items()}
        th[j] = math.radians(a)
        set_pose(arm, th)
        update()
        if CYL:
            note_stroke()
        hit = [h for h in robot_state(with_cell=False)]
        if hit:
            bad.append((round(a, 1), hit[:3]))
    rep["robot_sweep"][j] = dict(limits=[lo, hi], self_clash_angles=[b[0] for b in bad], first=bad[:2])
arm.animation_data.action = saved.get(arm.name)
update(1)
for ob in [body, rob] + [bpy.data.objects[n] for n in MOV]:
    rep["audit"][ob.name] = audit(ob)
ok = all(v["clash_steps"] == 0 for v in rep["travel"].values()) and rep["robot_cycle"]["clash_frames"] == 0 and \
    all(d["valid"] and d["simple"] for d in rep["drivers"].values())
if CYL:
    rep["j3_cylinder_stroke"] = STROKE
    ok = ok and STROKE["allowed_m"][0] <= STROKE["seen_m"][0] and STROKE["seen_m"][1] <= STROKE["allowed_m"][1]
rep["ok"] = ok
with open(os.path.join(HERE, "check_report%s.json" % ("_" + TAG if TAG else "")), "w") as fh:
    json.dump(rep, fh, indent=1)
print("CHECK OK" if ok else "CHECK CLASH", json.dumps(rep)[:6000])