"""Door robot arm: the rig in the ship scene (stage 'rig', Coder stage 2).

One arm per door and per version: DoorArm_AFT_A, DoorArm_AFT_B, DoorArm_FORE_A, DoorArm_FORE_B. Each is an object
hierarchy (dra_rigdata.LINKS: every moving part its own object, origin on its joint, rotation 0 and scale 1 at rest)
under a static root empty that hangs the arm under its pocket's rail (translation (0, -5.795, -4.182), turned 180 deg
about x). Nothing is keyed: every channel is driven, with simple expressions only, from the door's location (the doors
keep their own actions, frames 1-120) through a chain of solver empties (design_drivers.py, closed-form inverse
kinematics; DESIGN.md section 5):
  IK1 target (carriage c and the wrist centre in the arm frame) -> IK2 yaw -> IK3 shoulder -> IK4 elbow ->
  IK5 wrist matrix -> IK6 wrist angles -> the joints, the J3 cylinder and its piston, the carriage pinion.
Version B's wrist centre is 0.326 m off the door face (A: 0.63; CRITIQUE_1 item 1), so B solves its own pose.
Per door: a rail solver (c from the door) drives the energy chain's 44 links; the door fixtures follow the door.
The version switch: 'DoorArm | CONTROLS'["version"] (0 = A with the grab limbs, 1 = B the plate) drives hide_viewport
and hide_render of every object of each version; the collections 'DoorArm | version A' / '... B' hold them too.
The grip: CONTROLS["grip"] (0 open .. 1 closed on the knob; default 1) turns the four fingers of version A.
Cables over the elbow are Bezier curves parented to the lower arm whose first points hang on hooks to the upper arm and
whose middle points hang on empties blended between the two links (Copy Transforms, no Python).
"""
import math
import os
import sys
import copy
import json
import bpy
from mathutils import Vector, Matrix

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
if HERE not in sys.path:
    sys.path.insert(0, HERE)
import dra_rigdata as rd          # noqa: E402
import design_drivers as dd       # noqa: E402

def _d_t():
    with open(os.path.join(ROOT, "design.json")) as f:
        t = json.load(f)["tools"]
    return {"A": t["A"]["d_t"], "B": t["B"]["d_t"]}


D_T = _d_t()                      # wrist centre to the door face along the face normal (design.json tools; 2026-10-18)
DOORS = {"AFT": "Moonpool_Door_Aft", "FORE": "Moonpool_Door_Fore"}
CONTROLS = "DoorArm | CONTROLS"
GRIP = {"proximal": math.radians(14.7), "distal": math.radians(41.0)}
FINGER_CH = {"top": (1, 1.0), "bottom": (1, -1.0), "left": (2, -1.0), "right": (2, 1.0)}   # rotation index, sign
PIN_R = 0.040 * rd.S              # carriage pinion pitch radius (x the scale-up)


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


def coll(name, parent):
    c = bpy.data.collections.get(name) or bpy.data.collections.new(name)
    if c.name not in [x.name for x in parent.children]:
        parent.children.link(c)
    return c


def empty(name, c, size=0.08, kind='PLAIN_AXES', parent=None):
    o = bpy.data.objects.get(name)
    if o is None:
        o = bpy.data.objects.new(name, None)
        c.objects.link(o)
    o.empty_display_type = kind
    o.empty_display_size = size
    if parent is not None:
        o.parent = parent
        o.matrix_parent_inverse.identity()
    return o


def add_driver(ob, path, index, expr, vars_, report):
    """vars_: [(name, id object, data path)]"""
    if index is None:
        fc = ob.driver_add(path)
    else:
        fc = ob.driver_add(path, index)
    drv = fc.driver
    drv.type = 'SCRIPTED'
    for v in list(drv.variables):
        drv.variables.remove(v)
    for name, target, dpath in vars_:
        v = drv.variables.new()
        v.name = name
        v.type = 'SINGLE_PROP'
        v.targets[0].id_type = 'OBJECT'
        v.targets[0].id = target
        v.targets[0].data_path = dpath
    drv.expression = expr
    report.append((ob.name, path, index, fc))
    return fc


def apply_driver_list(lst, report, rename=None):
    """lst: design_drivers dicts; object names resolved in bpy.data.objects (rename maps design names to ours)"""
    rename = rename or {}
    for d in lst:
        tname = rename.get(d["target"], d["target"])
        ob = bpy.data.objects[tname]
        path = d["path"]
        if path.startswith('["'):
            key = path[2:-2]
            if key not in ob:
                ob[key] = 0.0
        vars_ = [(n, bpy.data.objects[rename.get(src, src)], p) for (n, src, p) in d["vars"]]
        add_driver(ob, path, d["index"], d["expr"], vars_, report)


# ------------------------------------------------------------------ the curves' hooks
def _hook(curve, name, target, point_ids, c_solver):
    """a Hook modifier moving whole Bezier points (handles too) with 'target'; bound at the rest pose"""
    m = curve.modifiers.new(name, 'HOOK')
    m.object = target
    idx = []
    for i in point_ids:
        idx += [3 * i, 3 * i + 1, 3 * i + 2]
    m.vertex_indices_set(idx)
    m.matrix_inverse = (curve.matrix_world.inverted() @ target.matrix_world).inverted()
    return m


def _blend_empty(name, c, a, b, infl_b, at_world):
    """an empty that sits halfway (by infl_b) between two link-carried empties a and b: Copy Transforms a (1.0) then b"""
    e = empty(name, c, 0.03)
    e.matrix_world = at_world
    k1 = e.constraints.new('COPY_TRANSFORMS')
    k1.target = a
    k1.influence = 1.0
    k2 = e.constraints.new('COPY_TRANSFORMS')
    k2.target = b
    k2.influence = infl_b
    return e


def hook_curves(prefix, objs, c_solver, c_arm):
    """cables over the elbow: points 0, 1 on the upper arm, 2 and 3 blended (0.35 / 0.70 toward the lower arm),
    4, 5 with the lower arm (the curve's parent). J3 hoses: points 0, 1 on the cylinder, 2 with the lower arm."""
    bpy.context.view_layer.update()
    ua, la = objs["upper arm"], objs["lower arm"]
    hu = empty("%s | hook upper" % prefix, c_solver, 0.05, parent=ua)
    hu.location = (0, 0, 0)
    hl = empty("%s | hook lower" % prefix, c_solver, 0.05, parent=la)
    hl.location = rd.J3S - rd.J2S
    bpy.context.view_layer.update()
    m1 = _blend_empty("%s | hook mid 1" % prefix, c_solver, hu, hl, 0.35, hu.matrix_world.copy())
    m2 = _blend_empty("%s | hook mid 2" % prefix, c_solver, hu, hl, 0.70, hu.matrix_world.copy())
    bpy.context.view_layer.update()
    for k in (1, 2):
        cu = bpy.data.objects.get("%s | cable elbow %d" % (prefix, k))
        if cu is None:
            continue
        _hook(cu, "Hook upper", hu, [0, 1], c_solver)
        _hook(cu, "Hook mid 1", m1, [2], c_solver)
        _hook(cu, "Hook mid 2", m2, [3], c_solver)
    cy = objs.get("J3 cylinder")
    for nm in ("rear", "front"):
        cu = bpy.data.objects.get("%s | J3 hose %s" % (prefix, nm))
        if cu is not None and cy is not None:
            _hook(cu, "Hook cylinder", cy, [0, 1], c_solver)


# ------------------------------------------------------------------ one arm
def build_arm(sf, key, ver, meshes, D, cols, report):
    prefix = "DoorArm_%s_%s" % (key, ver)
    c_arm = cols[ver]
    yb, zb = D["carriage"]["centre_y"], D["carriage"]["underside_z"]
    M_root = Matrix.Translation((0.0, yb, zb)) @ Matrix.Rotation(math.pi, 4, 'X')
    root, objs = rd.instantiate(c_arm, prefix, M_root, meshes, tools=(ver,))
    root["door"] = key
    root["version"] = ver
    for o in objs.values():
        o["version"] = ver
        o["door"] = key
    import dra_parts_arm
    dra_parts_arm.add_curves(c_arm, prefix, objs)
    for o in c_arm.objects:
        if o.name.startswith(prefix + " |") and o.type == 'CURVE':
            o["version"] = ver
            o["door"] = key
    hook_curves(prefix, objs, cols["solver"], c_arm)
    # solver empties and the IK / joint drivers (design_drivers.build with this version's wrist centre)
    Dv = copy.deepcopy(D)
    dr = Dv["doors"][key]
    att, nrm = dr["attach_local"], dr["normal_local"]
    dr["wc_local"]["A"] = [att[i] + D_T[ver] * nrm[i] for i in range(3)]
    if "tool_axes_C_" + ver in dr:                       # version A grips the knob turned (design.json tool_roll_A_deg)
        dr["tool_axes_C"] = dr["tool_axes_C_" + ver]
    plan = dd.build(Dv, key, prefix=prefix)
    for i, nm in enumerate(plan["solvers"]):
        e = empty(nm, cols["solver"], 0.05)
        e.location = (0, 0, 0)
    apply_driver_list(plan["drivers"], report)
    # the carriage pinion rolls on the rack: rotation about its axle = travel / pitch radius (in mesh at frame 1)
    c_closed = D["doors"][key]["carriage"]["c_closed"]
    pin = objs.get("carriage pinion")
    if pin is not None:
        add_driver(pin, "rotation_euler", 1, "(c - %.6f) / %.4f" % (c_closed, PIN_R),
                   [("c", bpy.data.objects["%s | IK1 target" % prefix], '["c"]')], report)
    # version A: the fingers on the grip control
    ctl = bpy.data.objects[CONTROLS]
    if ver == "A":
        for w, (ch, sg) in FINGER_CH.items():
            for part in ("proximal", "distal"):
                ob = objs.get("finger %s %s" % (w, part))
                if ob is None:
                    continue
                add_driver(ob, "rotation_euler", ch, "g * %.6f" % (sg * GRIP[part]), [("g", ctl, '["grip"]')], report)
    # visibility: the version switch
    expr = "v > 0.5" if ver == "A" else "v < 0.5"
    for o in [root] + list(objs.values()) + [o for o in c_arm.objects if o.name.startswith(prefix + " |") and o.type == 'CURVE']:
        for path in ("hide_viewport", "hide_render"):
            add_driver(o, path, None, expr, [("v", ctl, '["version"]')], report)
    return root, objs


# ------------------------------------------------------------------ per door: rail solver, chain, fixtures
def build_door_parts(sf, key, D, cols, report):
    door = bpy.data.objects[DOORS[key]]
    rs = empty("DoorArm_%s | rail solver" % key, cols["solver"], 0.05)
    rs["c"] = 0.0
    A_, B_ = D["doors"][key]["carriage"]["A"], D["doors"][key]["carriage"]["B"]
    add_driver(rs, '["c"]', None, "%.6f + %.6f*dx" % (A_, B_), [("dx", door, "location[0]")], report)
    D2 = copy.deepcopy(D)
    D2["energy_chain"][key]["moving_offset"] = (-0.020 if key == "AFT" else 0.020) * rd.S   # beside the bracket's drop plate
    sc_name, ch = dd.chain_drivers(D2, key, prefix="DoorArm_%s" % key)
    empty(sc_name, cols["solver"], 0.05)
    for d in ch:
        d["vars"] = [(n, ("DoorArm_%s | rail solver" % key) if src.endswith("| IK1 target") else src, p)
                     for (n, src, p) in d["vars"]]
    apply_driver_list(ch, report)
    # door fixtures follow the door (the door never turns): location = door location + attach offset
    att = D["doors"][key]["attach_local"]
    ctl = bpy.data.objects[CONTROLS]
    for ver in ("A", "B"):
        ob = bpy.data.objects.get("DoorArm_%s | door fixture %s" % (key, ver))
        if ob is None:
            continue
        for c in list(ob.users_collection):
            c.objects.unlink(ob)
        cols[ver].objects.link(ob)
        for i in range(3):
            add_driver(ob, "location", i, "d + %.6f" % att[i], [("d", door, "location[%d]" % i)], report)
        expr = "v > 0.5" if ver == "A" else "v < 0.5"
        for path in ("hide_viewport", "hide_render"):
            add_driver(ob, path, None, expr, [("v", ctl, '["version"]')], report)


def controls(c):
    ctl = empty(CONTROLS, c, 0.5, kind='CUBE')
    ctl.location = (3.05, -3.6, -3.0)
    ctl["version"] = 0
    ctl["grip"] = 1.0
    ui = ctl.id_properties_ui("version")
    ui.update(min=0, max=1, soft_min=0, soft_max=1, description="0 = A, with the grab limbs (gripper); 1 = B, the plate")
    ui = ctl.id_properties_ui("grip")
    ui.update(min=0.0, max=1.0, soft_min=0.0, soft_max=1.0, description="version A fingers: 0 open, 1 closed on the knob")
    return ctl


def build(sf, meshes):
    D = design()
    top = sf.collection
    cols = {"ctl": coll("DoorArm | controls", top), "A": coll("DoorArm | version A", top),
            "B": coll("DoorArm | version B", top), "solver": coll("DoorArm | solver (hidden)", top)}
    controls(cols["ctl"])
    report = []
    arms = {}
    for key in ("AFT", "FORE"):
        for ver in ("A", "B"):
            arms[(key, ver)] = build_arm(sf, key, ver, meshes, D, cols, report)
        build_door_parts(sf, key, D, cols, report)
    old = bpy.data.collections.get("DoorArm | door fixtures")
    if old is not None and len(old.objects) == 0:
        bpy.data.collections.remove(old)
    # the solver empties stay evaluated but hidden in the viewport (eye), never rendered
    for o in cols["solver"].objects:
        o.hide_render = True
    return arms, report


def finish_report(report):
    """after one evaluation: every driver simple and valid?"""
    bad = []
    for (name, path, index, fc) in report:
        drv = fc.driver
        if not (drv.is_simple_expression and drv.is_valid):
            bad.append([name, path, index, drv.expression, bool(drv.is_simple_expression), bool(drv.is_valid)])
    return dict(drivers=len(report), not_simple_or_invalid=bad)