"""The cell's rig, exec'd by rc_build.py after the objects exist (needs its globals: bpy, CAM, fk, set_pose, ROB, ...).

Sliders on RobotCell_Root (drivers are simple expressions: the user's Blender runs with Python auto-run off):
    door_left   0 shut .. 1 slid TRAVEL["L"] m to the right (the front track, in front of the right leaf)
    door_right  0 shut .. 1 slid TRAVEL["R"] m to the left (the back track, behind the left leaf)
    side_door   0 shut .. 1 swung out 100 degrees about its hinge
    turntable   the plate's angle in degrees (free)
    estop       0 out .. 1 pressed in 7 mm
The robot is an armature (RC_Robot_Rig, rc_robot.py): pose its bones A1..A6 and Tool_Spin directly; each turns about
one axis only (the others are locked) within the joint limits.

Timeline, 30 fps, frames 1-460 (the robot's action "Robot | WORK_CYCLE", the root's "RobotCell | DEMO"):
    1-30 the picture's pose; 60-115 down to the tray and up; 160-215 down to the turntable and up; 260 back to the
    picture's pose; 195-260 the turntable turns half a turn; 270-410 the front leaves slide open and shut; 280-420 the
    side door swings open and shut; 430-450 the emergency stop is pressed and let go.
Poses away from the picture are solved here by inverse kinematics (tool tip on a target, tool pointing down).
"""
import json, math
import numpy as np
from mathutils import Vector

root = bpy.data.objects["RobotCell_Root"]
arm = bpy.data.objects["RC_Robot_Rig"]
SLIDERS = {
    "door_left": (0.0, 0.0, 1.0, "0 shut, 1 slid open to the right (front track)"),
    "door_right": (0.0, 0.0, 1.0, "0 shut, 1 slid open to the left (back track)"),
    "side_door": (0.0, 0.0, 1.0, "0 shut, 1 swung out 100 degrees"),
    "turntable": (0.0, -3600.0, 3600.0, "the turntable plate's angle, degrees"),
    "estop": (0.0, 0.0, 1.0, "0 out, 1 pressed 7 mm"),
}
for k, (v, lo, hi, desc) in SLIDERS.items():
    root[k] = v
    root.id_properties_ui(k).update(min=lo, max=hi, soft_min=lo, soft_max=hi, description=desc)

RIG = {}


def drive(ob, path, idx, expr, prop, alt=None):
    ob.driver_remove(path, idx)
    fc = ob.driver_add(path, idx)
    d = fc.driver
    d.type = 'SCRIPTED'
    var = d.variables.new()
    var.name = "v"
    var.type = 'SINGLE_PROP'
    var.targets[0].id_type = 'OBJECT'
    var.targets[0].id = root
    var.targets[0].data_path = '["%s"]' % prop
    d.expression = expr
    if not d.is_simple_expression and alt:
        d.expression = alt
    for m in list(fc.modifiers):
        fc.modifiers.remove(m)
    d.expression = d.expression          # re-set, then mark valid (script-made drivers can stay invalid otherwise)
    d.is_valid = True
    RIG["%s.%s[%d]" % (ob.name, path, idx)] = dict(expression=d.expression, simple=d.is_simple_expression)


SS = "smoothstep(0.0, 1.0, v)"
SS_ALT = "min(max(v,0.0),1.0)*min(max(v,0.0),1.0)*(3-2*min(max(v,0.0),1.0))"
for key, prop in (("L", "door_left"), ("R", "door_right")):
    ob = bpy.data.objects["SM_RobotCell_Door%s" % ("Left" if key == "L" else "Right")]
    x0 = ob.location.x
    ob["closed_location"] = tuple(ob.location)
    ob["open_location"] = (x0 + TRAVEL[key], ob.location.y, ob.location.z)
    drive(ob, "location", 0, "%.6f + %.6f*%s" % (x0, TRAVEL[key], SS), prop, "%.6f + %.6f*%s" % (x0, TRAVEL[key], SS_ALT))
side = bpy.data.objects["SM_RobotCell_DoorSide"]
side["closed_rotation_deg"] = 0.0
drive(side, "rotation_euler", 2, "%.6f*%s" % (math.radians(100.0), SS), "side_door", "%.6f*%s" % (math.radians(100.0), SS_ALT))
plate = bpy.data.objects["SM_RobotCell_TurntablePlate"]
drive(plate, "rotation_euler", 2, "v*%.8f" % (math.pi / 180.0), "turntable")
es = bpy.data.objects["SM_RobotCell_EStop"]
es["closed_location"] = tuple(es.location)
drive(es, "location", 1, "%.6f + 0.007*min(max(v,0.0),1.0)" % es.location.y, "estop")

# ------------------------------------------------------------------ the robot's poses
fitp = os.path.join(HERE, "pose_fit.json")
PIC = json.load(open(fitp))["pose_deg"] if os.path.exists(fitp) else dict(A1=-22.0, A2=22.5, A3=3.5, A4=0.0, A5=40.6, A6=0.0)
KEYS = ["A1", "A2", "A3", "A5"]


def ik(tip_world, seed_deg, down=Vector((0, 0, -1)), it=300):
    """joint angles (degrees) that put the tool's tip on tip_world with the tool along `down` (A4 = A6 = 0)"""
    tgt = Vector(tip_world) - ROB
    q = np.array([math.radians(seed_deg[k]) for k in KEYS])

    def res(q):
        T = fk(dict(zip(KEYS, q)))
        tip = T["TIP"].translation
        ax = (T["TIP"].translation - T["FL"].translation).normalized()
        return np.array(list(tip - tgt) + list((ax - down) * 0.5))
    lam = 1e-3
    r = res(q)
    for _ in range(it):
        J = np.zeros((6, 4))
        for k in range(4):
            dq = np.zeros(4); dq[k] = 1e-5
            J[:, k] = (res(q + dq) - r) / 1e-5
        A, g = J.T @ J, J.T @ r
        st = np.linalg.solve(A + lam * np.eye(4), -g)
        r2 = res(q + st)
        if r2 @ r2 < r @ r:
            q, r = q + st, r2
            lam *= 0.5
        else:
            lam *= 5
        if r @ r < 1e-12:
            break
    out = {k: math.degrees(v) for k, v in zip(KEYS, q)}
    out.update(A4=0.0, A6=0.0)
    return out, float(np.sqrt(r @ r))


tx0, tx1, ty0, ty1 = TRAY
tray_c = Vector(((tx0 + tx1) / 2, (ty0 + ty1) / 2, Z_FLOOR + 0.004))
tt_c = Vector((TT.x - 0.30, TT.y - 0.04, TT_Z[1]))


def leaning(target, deg):
    """straight down, tipped `deg` away from the robot (a tool working at the far edge of its reach leans out)"""
    d = (target - ROB).to_2d().normalized()
    s = math.sin(math.radians(deg))
    return Vector((d.x * s, d.y * s, -math.cos(math.radians(deg))))


P_tray_up, e1 = ik(tray_c + Vector((0, 0, 0.36)), PIC)
P_tray, e2 = ik(tray_c + Vector((0, 0, 0.07)), P_tray_up)
P_tt_up, e3 = ik(tt_c + Vector((0, 0, 0.28)), P_tray_up, down=leaning(tt_c, 22))
P_tt, e4 = ik(tt_c + Vector((0, 0, 0.07)), P_tt_up, down=leaning(tt_c, 22))
for P in (P_tray_up, P_tray, P_tt_up, P_tt):
    P["Tool_Spin"] = 0.0
PICP = dict(PIC, Tool_Spin=0.0)
KEYF = [(1, PICP), (30, PICP), (60, P_tray_up), (80, P_tray), (95, dict(P_tray, Tool_Spin=720.0)),
        (115, dict(P_tray_up, Tool_Spin=720.0)), (160, dict(P_tt_up, Tool_Spin=720.0)), (180, dict(P_tt, Tool_Spin=720.0)),
        (195, dict(P_tt, Tool_Spin=1440.0)), (215, dict(P_tt_up, Tool_Spin=1440.0)), (260, dict(PICP, Tool_Spin=1440.0)),
        (460, dict(PICP, Tool_Spin=1440.0))]
for nm, (lo, hi) in LIMITS.items():
    for f, P in KEYF:
        assert lo <= P[nm] <= hi, (nm, f, P[nm])

sc = bpy.context.scene
sc.render.fps = 30
sc.frame_start, sc.frame_end = 1, 460
arm.animation_data_create()
for f, P in KEYF:
    set_pose(arm, {k: math.radians(v) for k, v in P.items()})
    for nm, ax in AXIS.items():
        arm.pose.bones[nm].keyframe_insert("rotation_euler", index="XYZ".index(ax), frame=f)
arm.animation_data.action.name = "Robot | WORK_CYCLE"
DEMO = {"turntable": [(1, 0.0), (195, 0.0), (260, 180.0)],
        "door_left": [(1, 0.0), (270, 0.0), (320, 1.0), (360, 1.0), (410, 0.0)],
        "door_right": [(1, 0.0), (270, 0.0), (300, 0.0)],
        "side_door": [(1, 0.0), (280, 0.0), (330, 1.0), (370, 1.0), (420, 0.0)],
        "estop": [(1, 0.0), (430, 0.0), (434, 1.0), (446, 1.0), (450, 0.0)]}
for prop, keys in DEMO.items():
    for f, v in keys:
        root[prop] = v
        root.keyframe_insert('["%s"]' % prop, frame=f)
root.animation_data.action.name = "RobotCell | DEMO"
for k in SLIDERS:
    root[k] = 0.0
sc.frame_set(1)
root["rig"] = json.dumps(RIG)
arm["poses_deg"] = json.dumps(dict(picture=PIC, tray_up=P_tray_up, tray=P_tray, turntable_up=P_tt_up, turntable=P_tt))
arm["ik_error_m"] = json.dumps([round(e, 5) for e in (e1, e2, e3, e4)])