"""Driver plan as data: every driver the rig needs, as Blender SIMPLE expressions (numbers, variables, + - * /,
comparisons, 'a if c else b', min max clamp abs sqrt sin cos atan2 acos ...; no ** and no Python).
design_make_json.py stores the result in design.json["drivers"]; design_driver_proto.py builds it in Blender to prove it.

A driver = dict(target=object name, path=data path, index=int or None, expr=string,
                vars=[(name, source object name, data path)]).
Solver empties carry custom properties; each empty only reads EARLIER empties (no property reads its own object).
"""
import math


def g(v):
    return "%.10g" % v


def build(design, key, prefix=None):
    D = design
    K = D["kinematics"]
    dr = D["doors"][key]
    P = prefix or ("DoorArm_%s" % key)
    door = dr["door"]
    tool = D["tools"]["A"]                       # both versions share the wrist centre
    wc = dr["wc_local"]["A"]
    yb, zb = D["carriage"]["centre_y"], D["carriage"]["underside_z"]
    A, B = dr["carriage"]["A"], dr["carriage"]["B"]
    j2r, j2h = K["j2_r"], K["j2_h"]
    L2, L3, beta, phi_l0 = K["L2"], K["L3"], K["beta_rad"], K["phi_l0_rad"]
    xt, yt, zt = dr["tool_axes_C"]["x_t"], dr["tool_axes_C"]["y_t"], dr["tool_axes_C"]["z_t"]
    rod = K["rod"]
    tl = rod["trunnion_in_lower"]
    j3l = [K["joints"]["J3"]["origin"][i] - K["joints"]["J2"]["origin"][i] for i in range(3)]
    ka, kb = rod["knob_in_upper"][0], rod["knob_in_upper"][2]
    d0x = rod["knob_rest_A"][0] - rod["trunnion_rest_A"][0]
    d0z = rod["knob_rest_A"][2] - rod["trunnion_rest_A"][2]
    L0 = math.hypot(d0x, d0z)
    b0 = math.atan2(d0z, d0x)
    ux, uz = d0x / L0, d0z / L0
    S1, S2, S3, S4, S5, S6 = ["%s | IK%d %s" % (P, i, n) for i, n in
                              ((1, "target"), (2, "yaw"), (3, "shoulder"), (4, "elbow"), (5, "wrist matrix"), (6, "wrist angles"))]
    drv = []

    def add(target, path, expr, vars_, index=None):
        drv.append(dict(target=target, path=path, index=index, expr=expr, vars=vars_))

    door_vars = [("dx", door, "location[0]"), ("dy", door, "location[1]"), ("dz", door, "location[2]")]
    # IK1: carriage and the wrist centre in the hanging arm frame C (C = world turned 180 deg about x)
    add(S1, '["c"]', "%s + %s*dx" % (g(A), g(B)), door_vars[:1])
    add(S1, '["X"]', "%s*dx + %s" % (g(1 - B), g(wc[0] - A)), door_vars[:1])
    add(S1, '["Y"]', "%s - dy" % g(yb - wc[1]), door_vars[1:2])
    add(S1, '["H"]', "%s - dz" % g(zb - wc[2]), door_vars[2:3])
    v1 = lambda n: (n, S1, '["%s"]' % n)
    # IK2: turret yaw; the wrist centre seen from J2 in the arm's plane (R along, Hh up)
    add(S2, '["q1"]', "atan2(Y, X)", [v1("X"), v1("Y")])
    add(S2, '["R"]', "sqrt(X*X + Y*Y) + %s" % g(j2r), [v1("X"), v1("Y")])
    add(S2, '["Hh"]', "H - %s" % g(j2h), [v1("H")])
    v2 = lambda n: (n, S2, '["%s"]' % n)
    # IK3: lower arm direction (elbow-up as in the reference) and q2
    phil = "atan2(Hh, R) + acos(clamp((%s + R*R + Hh*Hh) / (%s*sqrt(R*R + Hh*Hh)), -1, 1))" % (g(L2 * L2 - L3 * L3), g(2 * L2))
    add(S3, '["phi_l"]', phil, [v2("R"), v2("Hh")])
    add(S3, '["q2"]', "%s - %s" % (phil, g(phi_l0)), [v2("R"), v2("Hh")])
    # IK4: forearm direction and q3
    phif = "atan2(Hh - %s*sin(phi_l), R - %s*cos(phi_l)) - %s" % (g(L2), g(L2), g(beta))
    v4 = [v2("R"), v2("Hh"), ("phi_l", S3, '["phi_l"]')]
    add(S4, '["phi_f"]', phif, v4)
    add(S4, '["q3"]', "%s - phi_l + %s" % (phif, g(phi_l0)), v4)
    # IK5: wrist matrix entries M = R3^T R_tool (the five used) and the knob seen from the trunnion (lower-arm frame)
    Pv = lambda v: "(cos(q1)*%s + sin(q1)*%s)" % (g(v[0]), g(v[1]))
    v5 = [("q1", S2, '["q1"]'), ("phi_f", S4, '["phi_f"]')]
    add(S5, '["M11"]', "cos(phi_f)*%s + sin(phi_f)*%s" % (Pv(xt), g(xt[2])), v5)
    add(S5, '["M21"]', "-sin(q1)*%s + cos(q1)*%s" % (g(xt[0]), g(xt[1])), v5[:1])
    add(S5, '["M31"]', "cos(phi_f)*%s - sin(phi_f)*%s" % (g(xt[2]), Pv(xt)), v5)
    add(S5, '["M12"]', "cos(phi_f)*%s + sin(phi_f)*%s" % (Pv(yt), g(yt[2])), v5)
    add(S5, '["M13"]', "cos(phi_f)*%s + sin(phi_f)*%s" % (Pv(zt), g(zt[2])), v5)
    vq3 = [("q3", S4, '["q3"]')]
    add(S5, '["kx"]', "%s + %s*cos(q3) - %s*sin(q3)" % (g(j3l[0] - tl[0]), g(ka), g(kb)), vq3)
    add(S5, '["kz"]', "%s + %s*sin(q3) + %s*cos(q3)" % (g(j3l[2] - tl[2]), g(ka), g(kb)), vq3)
    # IK6: wrist angles, branch q5 < 0 (the tool bends toward the forearm's +Z side, where the wrist block is cut back);
    # the J3 cylinder's length and its turn in the lower arm
    m = lambda n: (n, S5, '["%s"]' % n)
    add(S6, '["q5"]', "-acos(clamp(M11, -1, 1))", [m("M11")])
    add(S6, '["q4"]', "atan2(-M21, M31)", [m("M21"), m("M31")])
    add(S6, '["q6"]', "atan2(-M12, -M13)", [m("M12"), m("M13")])
    add(S6, '["L"]', "sqrt(kx*kx + kz*kz)", [m("kx"), m("kz")])
    add(S6, '["cyl"]', "atan2(kz, kx) - %s" % g(b0), [m("kx"), m("kz")])
    # joints (object names as in DESIGN.md)
    O = lambda n: "%s | %s" % (P, n)
    add(O("carriage"), "location", "c", [("c", S1, '["c"]')], 0)
    add(O("turret"), "rotation_euler", "q1", [("q1", S2, '["q1"]')], 2)
    add(O("lower arm"), "rotation_euler", "-q2", [("q2", S3, '["q2"]')], 1)
    add(O("upper arm"), "rotation_euler", "-q3", [("q3", S4, '["q3"]')], 1)
    add(O("forearm"), "rotation_euler", "q4", [("q4", S6, '["q4"]')], 0)
    add(O("wrist head"), "rotation_euler", "q5", [("q5", S6, '["q5"]')], 1)
    add(O("flange"), "rotation_euler", "q6", [("q6", S6, '["q6"]')], 0)
    add(O("J3 cylinder"), "rotation_euler", "-cyl", [("cyl", S6, '["cyl"]')], 1)
    add(O("J3 piston"), "location", "%s + %s*(L - %s)" % (g(d0x), g(ux), g(L0)), [("L", S6, '["L"]')], 0)
    add(O("J3 piston"), "location", "%s + %s*(L - %s)" % (g(d0z), g(uz), g(L0)), [("L", S6, '["L"]')], 2)
    solvers = [S1, S2, S3, S4, S5, S6]
    for d in drv:
        assert len(d["expr"]) <= 250, (d["target"], d["path"], len(d["expr"]))
    return dict(solvers=solvers, drivers=drv, piston_rest=[d0x, 0.0, d0z], cyl_beta0=b0, rod_L0=L0, rod_axis=[ux, 0.0, uz])


def chain_drivers(design, key, prefix=None):
    """energy chain: a solver empty (moving end M, moving-run length a, bend start bx) from the carriage x, then per link
    (0 at the carriage end) location x, y and rotation z. Returns (solver name, drivers)."""
    D = design
    ch = D["energy_chain"][key]
    P = prefix or ("DoorArm_%s" % key)
    S1 = "%s | IK1 target" % P
    SC = "%s | chain solver" % P
    F, Lc, R, p = ch["fixed_x"], ch["length"], ch["bend_radius"], ch["pitch"]
    ym, yf = ch["y_moving_run"], ch["y_fixed_run"]
    sgn = ch["bend_toward"]          # -1: the loop lies toward -x of both ends (AFT); +1 (FORE)
    yc = (ym + yf) / 2
    pr = math.pi * R
    out = []
    vc = [("c", S1, '["c"]')]
    out.append(dict(target=SC, path='["M"]', index=None, expr="c + %s" % g(ch["moving_offset"]), vars=vc))
    out.append(dict(target=SC, path='["a"]', index=None, expr="(%s - %s*c + %s) / 2" % (g(sgn * (F - ch["moving_offset"])), g(sgn), g(Lc - pr)), vars=vc))
    out.append(dict(target=SC, path='["bx"]', index=None, expr="c + %s + %s*(%s - %s*c + %s) / 2" % (
        g(ch["moving_offset"]), g(sgn), g(sgn * (F - ch["moving_offset"])), g(sgn), g(Lc - pr)), vars=vc))
    k = sgn * (1 if ym < yf else -1)
    base = math.pi if sgn < 0 else 0.0
    vs = [("M", SC, '["M"]'), ("a", SC, '["a"]'), ("bx", SC, '["bx"]')]
    for i in range(ch["links"]):
        s = (i + 0.5) * p
        th = "clamp((%s - a) / %s, 0, %s)" % (g(s), g(R), g(math.pi))
        x_expr = "M + %s if %s < a else (bx + %s*sin(%s) if %s < a + %s else bx + %s*(%s - a - %s))" % (
            g(sgn * s), g(s), g(sgn * R), th, g(s), g(pr), g(-sgn), g(s), g(pr))
        y_expr = "%s + %s*cos(%s)" % (g(yc), g(ym - yc), th)
        r_expr = "%s + %s*%s" % (g(base), g(k), th)
        name = "%s | chain link %02d" % (P, i)
        for path, idx, e in (("location", 0, x_expr), ("location", 1, y_expr), ("rotation_euler", 2, r_expr)):
            assert len(e) <= 250, (name, len(e))
            out.append(dict(target=name, path=path, index=idx, expr=e, vars=vs))
    return SC, out