"""Door robot arm: the showroom (both versions upright side by side on a floor rail) and the parts the design adds in
the ship scene: the rail kits (running flanges, rack, end stops, energy-chain tray), the chain links and the door
fixtures. Ship coordinates (metres); FORE is the AFT kit mirrored in the plane x = 3.05.

Showroom height: the design put the rail at z 0.25..0.585 with the carriage on top of it; but the carriage wraps the
rail's top 0.255 m (as it wraps the rail's lower 0.255 m in the ship), so the rail is at z 0.505..0.840 here
(turntable face stays at z 1.012). Arm B stands a whole number of rack teeth from arm A (x +1.8946, not +1.9), so both
pinions mesh with the same rack.
"""
import math
import json
import os
import bmesh
import bpy
from mathutils import Vector, Matrix
from mathutils.bvhtree import BVHTree
from dra_lib import (Link, lathe, slab, rbox, tube, merge, xform, frame, basis, mirror_y, W, G, S, R, O, MAT_NAMES, uv_tile)
from dra_shapes import rrect, xz_slab, yz_slab, xy_slab, screw, dome, hexbolt, cyl, ring
import dra_rigdata as rd
from dra_parts_base import RACK_PITCH, PIN_MOD, PIN_Z

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
ROOT_Z = 1.012                     # showroom turntable face height
ARM_X = {"A": -1.9, "B": -1.9 + 151 * RACK_PITCH}
MIRROR_X = 3.05


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


def link_collection(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 mesh_from_link(L, name, origin=(0, 0, 0)):
    import dra_lib
    me = dra_lib.finish(L, name, origin=origin)
    uv_tile(me)
    if "UVMap" not in me.uv_layers:
        uv = me.uv_layers.new(name="UVMap")
        # a plain planar map; the bake stage repacks every mesh into the atlas
        for li, l in enumerate(me.loops):
            co = me.vertices[l.vertex_index].co
            uv.data[li].uv = (co.x * 0.25 + 0.5, co.y * 0.25 + 0.5)
        me.uv_layers.active = uv
        uv.active_render = True
    return me


# ------------------------------------------------------------------ rack (teeth along +local Z)
def rack_piece(x0, x1, phase, zb, zt, zr, half_w=0.030, mat=S):
    """a gear rack from x0 to x1 (local), body from zb to zr (root line), teeth from zr to zt, tooth centres at
    phase + k * pitch. Teeth point along +Z."""
    L = Link("rack")
    L.put(rbox((x1 - x0, 2 * half_w, zr - zb), 0.002), Matrix.Translation(((x0 + x1) / 2, 0, (zb + zr) / 2)), mat, 'rack')
    p = RACK_PITCH
    k0 = math.ceil((x0 - phase) / p + 0.5)
    k1 = math.floor((x1 - phase) / p - 0.5)
    for k in range(k0, k1 + 1):
        cx = phase + k * p
        prof = [(cx - 0.0068, zr - 0.001, 0.0), (cx + 0.0068, zr - 0.001, 0.0), (cx + 0.0028, zt, 0.0), (cx - 0.0028, zt, 0.0)]
        L.put(xz_slab(prof, -half_w + 0.001, half_w - 0.001, 0.0, 0.0, n=1), None, mat, 'rack t', seats=('rack',))
    return L


# ------------------------------------------------------------------ showroom
def build_showroom(sc, meshes):
    top = sc.collection
    c_arms = link_collection("Showroom | arms", top)
    c_studio = link_collection("Showroom | studio", top)
    for v in ("A", "B"):
        c = link_collection("Showroom | arm " + v, c_arms)
        M = Matrix.Translation((ARM_X[v], 0.0, ROOT_Z))
        root, objs = rd.instantiate(c, "Showroom | arm " + v, M, meshes, tools=(v,))
        root["version"] = v
        for ob in objs.values():
            ob["version"] = v
        # the cables over the elbow and the cylinder hoses (curves at rest)
        import dra_parts_arm
        dra_parts_arm.add_curves(c, "Showroom | arm " + v, objs)
        for ob in c.objects:
            ob["version"] = v
    _floor_rail(c_studio)
    _studio(sc, c_studio)


def _floor_rail(c):
    """the ship's rail turned over on short posts: box 0.458 x 0.335, running flanges and the rack on its top"""
    L = Link("floor rail")
    x0, x1 = -3.6, 3.6
    zt = ROOT_Z - 0.172                   # rail top (= its bottom face in the ship)
    zb = zt - 0.335
    L.put(rbox((x1 - x0, 0.458, 0.335), 0.004), Matrix.Translation((0, 0, (zt + zb) / 2)), G, 'rail')
    for s in (1, -1):
        L.put(rbox((x1 - x0, 0.050, 0.030), 0.0025), Matrix.Translation((0, s * 0.254, zt - 0.015)), S, 'K01', seats=('rail',))
        for k in range(int((x1 - x0) / 0.25)):
            x = x0 + 0.125 + k * 0.25
            L.put(dome(0.009, 0.004, S), frame((x, s * 0.279, zt - 0.015), (0, s, 0)), None, 'K01b', seats=('K01',))
    # rack: in carriage terms body Z -0.172..-0.150, teeth to -0.132; teeth between the two arms' pinion teeth
    rk = rack_piece(x0 + 0.05, x1 - 0.05, ARM_X["A"] + RACK_PITCH / 2, zt, ROOT_Z - 0.132, ROOT_Z - 0.150)
    merge(L.bm, rk.bm)
    for p in rk.parts:
        L.pid(p['id'], p['seats'])
    rk.bm.free()
    # posts: square tubes 0.12 x 0.12 every 2.2 m
    for x in (-3.3, -1.1, 1.1, 3.3):
        if abs(x - ARM_X["A"]) < 0.5 or abs(x - ARM_X["B"]) < 0.5:
            pass
        L.put(rbox((0.12, 0.12, zb), 0.006), Matrix.Translation((x, 0, zb / 2)), G, 'post', seats=('rail',))
        L.put(xy_slab(rrect(0.22, 0.22, 0.01, x, 0), 0.0, 0.012, 0.0, 0.003), None, G, 'post', seats=('post',))
    me = mesh_from_link(L, "Showroom | floor rail")
    ob = bpy.data.objects.new("Showroom | floor rail", me)
    c.objects.link(ob)
    return ob


def _studio(sc, c):
    m = bpy.data.materials.get("M_Studio_Floor") or bpy.data.materials.new("M_Studio_Floor")
    m.diffuse_color = (0.55, 0.55, 0.55, 1)
    m.use_nodes = True
    m.node_tree.nodes["Principled BSDF"].inputs["Base Color"].default_value = (0.55, 0.55, 0.55, 1)
    m.node_tree.nodes["Principled BSDF"].inputs["Roughness"].default_value = 0.6
    # cyclorama: floor from y +12 to -4, a 6 m cove up to z 9 behind the arms (-y), seen from the left side (+y)
    bm = bmesh.new()
    prof = [(12.0, 0.0), (-4.0, 0.0)]
    for k in range(1, 13):
        a = (math.pi / 2) * k / 12
        prof.append((-4.0 - 6.0 * math.sin(a), 6.0 - 6.0 * math.cos(a)))
    prof.append((-10.0, 9.0))
    xs = [-14.0 + 2.0 * i for i in range(15)]
    grid = [[bm.verts.new((x, y, z)) for x in xs] for (y, z) in prof]
    for j in range(len(prof) - 1):
        for i in range(len(xs) - 1):
            f = bm.faces.new((grid[j][i], grid[j][i + 1], grid[j + 1][i + 1], grid[j + 1][i]))
    bmesh.ops.recalc_face_normals(bm, faces=bm.faces[:])
    for f in bm.faces:
        if f.normal.z < -0.1 or (f.normal.y < -0.1):
            f.normal_flip()
    me = bpy.data.meshes.new("Showroom | cyclorama")
    bm.to_mesh(me)
    bm.free()
    me.materials.append(m)
    me.shade_smooth()
    uv_tile(me, 4.0)
    ob = bpy.data.objects.new("Showroom | cyclorama", me)
    c.objects.link(ob)
    for nm, loc, size, energy, rot in (("key", (-5.0, 6.0, 7.0), (2.0, 2.0), 800, None),
                                         ("fill", (6.0, 5.0, 3.0), (3.0, 3.0), 300, None),
                                         ("rim", (0.0, -7.0, 5.0), (1.0, 4.0), 500, None)):
        ld = bpy.data.lights.new("Showroom | " + nm, 'AREA')
        ld.energy = energy
        ld.shape = 'RECTANGLE'
        ld.size, ld.size_y = size
        lo = bpy.data.objects.new("Showroom | light " + nm, ld)
        lo.location = loc
        d = (Vector((0, 0, 1.6)) - Vector(loc)).normalized()
        lo.rotation_euler = d.to_track_quat('-Z', 'Y').to_euler()
        c.objects.link(lo)
    w = bpy.data.worlds.get("Showroom | world") or bpy.data.worlds.new("Showroom | world")
    w.color = (0.35, 0.35, 0.35)
    w.use_nodes = True
    bg = w.node_tree.nodes.get("Background")
    if bg:
        bg.inputs[0].default_value = (0.35, 0.35, 0.35, 1)
        bg.inputs[1].default_value = 0.4
    sc.world = w


# ------------------------------------------------------------------ ship fit: rail kits, chain, fixtures
def _mx(side):
    """AFT: identity; FORE: mirror in the plane x = 3.05"""
    if side == "AFT":
        return Matrix.Identity(4)
    return Matrix.Translation((2 * MIRROR_X, 0, 0)) @ Matrix.Scale(-1, 4, Vector((1, 0, 0)))


def rail_kit_link(d):
    """the AFT rail kit in ship coordinates"""
    sc = d["scene"]
    xr0, xr1 = sc["rail"]["usable_AFT"]
    y0, y1 = sc["rail"]["y"]
    zb, zt = sc["rail"]["z"]
    yc = (y0 + y1) / 2
    L = Link("rail kit")
    # K01 running flanges on both side faces at the rail's bottom, low dome bolt heads on their outer faces
    for side, yy in ((-1, y0 - 0.025), (1, y1 + 0.025)):
        L.put(rbox((xr1 - xr0, 0.050, 0.030), 0.0025), Matrix.Translation(((xr0 + xr1) / 2, yy, zb + 0.015)), S, 'K01')
        x = xr0 + 0.125
        while x < xr1 - 0.05:
            L.put(dome(0.009, 0.004, S), frame((x, yy + side * 0.025, zb + 0.015), (0, side, 0)), None, 'K01b', seats=('K01',))
            x += 0.25
    # K02 rack under the rail, teeth down; tooth gaps where the pinion's bottom tooth sits at the closed door (frame 1)
    c_closed = d["carriage"]["travel_ends"]["AFT"][1]
    rk = rack_piece(-8.225, -2.905, c_closed + RACK_PITCH / 2, 0.0, 0.040, 0.022, 0.030)
    xform(rk.bm, Matrix.Translation((0, yc, zb)) @ Matrix.Scale(-1, 4, Vector((0, 0, 1))))
    lay_r = rk.bm.faces.layers.int.get('part')
    remap = {i: L.pid(p['id'], p['seats']) for i, p in enumerate(rk.parts)}
    for f in rk.bm.faces:
        f[lay_r] = remap.get(f[lay_r], f[lay_r])
    merge(L.bm, rk.bm)
    rk.bm.free()
    # bolts holding the rack to the rail's underside (between the teeth rows, visible from below on the body's flanks)
    # K03 end stops: steel block + rubber buffer, under the rail between the flanges
    for (bx0, bx1, rx0, rx1) in ((-8.685, -8.585, -8.585, -8.555), (-2.545, -2.475, -2.575, -2.545)):
        L.put(rbox((bx1 - bx0, 0.30, 0.12), 0.004), Matrix.Translation(((bx0 + bx1) / 2, yc, -4.07)), G, 'K03')
        L.put(rbox((rx1 - rx0 + 0.002, 0.26, 0.10), 0.006), Matrix.Translation(((rx0 + rx1) / 2, yc, -4.07)), R, 'K03r',
              seats=('K03',))
        for yy in (-0.1, 0.1):
            L.put(hexbolt(0.018, 0.008), frame(((bx0 + bx1) / 2, yc + yy, -4.13), '-Z'), None, 'K03b', seats=('K03',))
    # the inner end stop's block also bolts to the rail's end face: an angle plate up the end face
    L.put(rbox((0.012, 0.30, 0.20), 0.003), Matrix.Translation((-2.539, yc, -3.92)), G, 'K03', seats=('K03',))
    # K04 energy-chain tray: floor, two walls, end walls, the slanted well-side end; hangers on the far wall
    t = d["energy_chain"]["AFT"]["tray"]
    tx0, tx1s = t["x"]
    ty0, ty1 = t["y"]
    fz = t["floor_top_z"]
    wz = t["wall_top_z"]
    xa, xb = t["end_slant"]["x"]              # well-side end x at y0 and at y1
    floor = [(tx0, ty0), (xa, ty0), (xb, ty1), (tx0, ty1)]
    L.put(xy_slab(floor, fz - 0.006, fz, 0.0015, 0.0015), None, G, 'K04')
    for yy, nm in ((ty0, 'near'), (ty1, 'far')):
        xe = xa if yy == ty0 else xb
        L.put(rbox((xe - tx0, 0.006, wz - fz + 0.006), 0.0015), Matrix.Translation(((tx0 + xe) / 2, yy + (0.003 if yy == ty0 else -0.003),
                                                                                    (wz + fz - 0.006) / 2)), G, 'K04', seats=('K04',))
    ang = math.atan2(ty1 - ty0, xb - xa)
    ln = math.hypot(xb - xa, ty1 - ty0)
    Mend = Matrix.Translation(((xa + xb) / 2, (ty0 + ty1) / 2, (wz + fz - 0.006) / 2)) @ Matrix.Rotation(ang, 4, 'Z')
    L.put(rbox((ln - 0.004, 0.006, wz - fz + 0.005), 0.0015), Mend, G, 'K04e', seats=('K04',))
    L.put(rbox((0.006, ty1 - ty0 - 0.004, wz - fz + 0.005), 0.0015),
          Matrix.Translation((tx0 + 0.0035, (ty0 + ty1) / 2, (wz + fz - 0.006) / 2 + 0.0005)), G, 'K04e', seats=('K04',))
    x = tx0 + 0.3
    deck = sc["deck_underside_z"]
    while x < xb - 0.2:
        L.put(rbox((0.04, 0.006, deck - wz - 0.0035), 0.001), Matrix.Translation((x, ty1 + 0.003, (deck + wz) / 2 - 0.0028)), G, 'K04h',
              seats=('K04',))
        L.put(rbox((0.06, 0.06, 0.006), 0.001), Matrix.Translation((x, ty1 + 0.03, deck - 0.0035)), G, 'K04h', seats=('K04h',))
        x += 0.8
    # the chain's fixed end bracket on the tray floor at the fixed x
    fy = d["energy_chain"]["AFT"]["y_fixed_run"]
    c1 = d["carriage"]["travel_ends"]["AFT"][1]
    pos = chain_positions(d, "AFT", c1)
    (x43, y43, r43), (x42, y42, r42) = pos[-1], pos[-2]
    sg = 1.0 if x43 > x42 else -1.0
    xe = x43 + sg * 0.0505
    x0b, x1b = (xe, xe + 0.012) if sg > 0 else (xe - 0.012, xe)
    L.put(yz_slab(rrect(0.07, 0.10, 0.006, fy, (fz + wz) / 2 - 0.005), x0b, x1b, 0.0015, 0.0015), None, S, 'K05f',
          seats=('K04',))
    return L


def chain_link_mesh():
    """one energy-chain link lying on its side (bend axis vertical), centre at the origin, x along the chain (pitch
    0.075). Outer plates (top and bottom) on the back half, inner plates one step further in on the front half, so a
    link's inner plates slide under the next link's outer plates; the bars stand outside the inner plates' width.
    Kept light (it is repeated 88 times): plain chamfered boxes."""
    L = Link("chain link")
    for sz in (1, -1):
        M = Matrix.Scale(sz, 4, Vector((0, 0, 1)))
        L.put(rbox((0.054, 0.068, 0.006), 0.0012, segs=1), M @ Matrix.Translation((-0.023, 0, 0.047)), R, 'outer')
        L.put(rbox((0.050, 0.054, 0.006), 0.0012, segs=1), M @ Matrix.Translation((0.021, 0, 0.039)), R, 'inner', seats=('step',))
        L.put(rbox((0.008, 0.054, 0.0145), 0.001, segs=1), M @ Matrix.Translation((0.0, 0, 0.0428)), R, 'step', seats=('outer', 'inner'))
    for yy in (-0.031, 0.031):
        L.put(rbox((0.016, 0.006, 0.0875), 0.0012, segs=1), Matrix.Translation((-0.006, yy, 0)), R, 'bar', seats=('outer',))
    me = mesh_from_link(L, "DoorArm | chain link")
    return me


def chain_positions(d, side, c):
    """the designer's chain driver expressions evaluated for carriage position c -> [(x, y, rot_z)] (ship, AFT or FORE)"""
    dr = d["drivers"][side]["chain_link_drivers"]
    env = {"sin": math.sin, "cos": math.cos, "clamp": lambda v, a, b: max(a, min(b, v)), "pi": math.pi,
           "min": min, "max": max}
    solver = {}
    vals = {}
    for e in dr:
        if e["target"].endswith("chain solver"):
            loc = dict(env)
            loc.update({v[0]: c for v in e["vars"]})
            solver[e["path"].strip('[]"')] = eval(e["expr"], {"__builtins__": {}}, loc)
    for e in dr:
        if e["target"].endswith("chain solver"):
            continue
        loc = dict(env)
        for v in e["vars"]:
            key = v[2].strip('[]"')
            loc[v[0]] = solver[key]
        val = eval(e["expr"], {"__builtins__": {}}, loc)
        vals.setdefault(e["target"], {})[(e["path"], e["index"])] = val
    out = []
    for i in range(d["energy_chain"][side]["links"]):
        nm = "DoorArm_%s | chain link %02d" % (side, i)
        v = vals.get(nm, {})
        out.append((v.get(("location", 0), 0.0), v.get(("location", 1), 0.0), v.get(("rotation_euler", 2), 0.0)))
    return out


def fixture_links(d, door_ob, side):
    """door fixtures A (pad, base plate, stem, mushroom knob) and B (conformal pad) for one door, in ship coordinates
    relative to the attach point (rotation 0: the mesh is built in the ship's axes)"""
    dd = d["doors"][side]
    loc = Vector(dd["closed_location"])
    att = loc + Vector(dd["attach_local"])
    n = Vector(dd["normal_local"]).normalized()
    xa = Vector((1, 0, 0))
    ya = n.cross(xa).normalized()
    Mf = Matrix((xa, ya, n)).transposed().to_4x4()          # fixture local -> ship axes
    me = door_ob.data
    verts = [loc + v.co for v in me.vertices]
    bvh = BVHTree.FromPolygons(verts, [tuple(p.vertices) for p in me.polygons])
    out = {}
    for ver, size, top in (("A", 0.32, 0.020), ("B", 0.60, 0.020)):
        L = Link("door fixture " + ver)
        N = 8 if ver == "A" else 12
        bm = bmesh.new()
        topv, botv = [], []
        for j in range(N + 1):
            rt, rb = [], []
            for i in range(N + 1):
                u = -size / 2 + size * i / N
                v = -size / 2 + size * j / N
                p = att + xa * u + ya * v
                hit = bvh.ray_cast(p + n * 0.3, -n, 1.0)
                hp = hit[0] if hit[0] is not None else p
                lp = Mf.inverted() @ (hp - att).to_4d()
                q = Vector((u, v, lp.z + 0.0005))
                rb.append(bm.verts.new(Mf.to_3x3() @ q))
                rt.append(bm.verts.new(Mf.to_3x3() @ Vector((u, v, top))))
            topv.append(rt)
            botv.append(rb)
        from dra_lib import face as _face
        for j in range(N):
            for i in range(N):
                _face(bm, [topv[j][i], topv[j][i + 1], topv[j + 1][i + 1], topv[j + 1][i]], G, out=n)
                _face(bm, [botv[j][i], botv[j][i + 1], botv[j + 1][i + 1], botv[j + 1][i]], G, out=-n)
        per = [(i, 0) for i in range(N)] + [(N, j) for j in range(N)] + [(i, N) for i in range(N, 0, -1)] + [(0, j) for j in range(N, 0, -1)]
        for k in range(len(per)):
            (i0, j0), (i1, j1) = per[k], per[(k + 1) % len(per)]
            a, b = topv[j0][i0], topv[j1][i1]
            c_, e = botv[j1][i1], botv[j0][i0]
            mid = (a.co + b.co) / 2
            _face(bm, [a, b, c_, e], G, out=mid - n * mid.dot(n))
        L.put(bm, None, G, 'D%s1' % ver)
        R3 = Mf
        if ver == "A":
            L.put(rbox((0.26, 0.26, 0.030), 0.004, rv=0.02, vax=2), R3 @ Matrix.Translation((0, 0, top + 0.015)), G, 'DA2', seats=('DA1',))
            for (u, v) in ((0.095, 0.095), (-0.095, 0.095), (0.095, -0.095), (-0.095, -0.095)):
                L.put(hexbolt(0.024, 0.010), R3 @ Matrix.Translation((u, v, top + 0.030)), None, 'DA2b', seats=('DA2',))
            L.put(cyl(0.030, 0.115, 32, 0.003, S, z0=top + 0.030), R3, S, 'DA3', seats=('DA2', 'DA4'))
            L.put(lathe([(0.0, 0.165, 0, S), (0.065, 0.165, 0.003, S), (0.065, 0.215, 0.004, S), (0.0, 0.215, 0, S)], 48), R3, S,
                  'DA4', seats=('DA3',))
        out[ver] = (L, att)
    return out


def build_ship_fit(sc, meshes):
    d = _design()
    top = sc.collection
    c_rail = link_collection("DoorArm | rail kit", top)
    c_fix = link_collection("DoorArm | door fixtures", top)
    L = rail_kit_link(d)
    me_aft = mesh_from_link(L, "DoorArm | rail kit AFT")
    chain_me = chain_link_mesh()
    for side in ("AFT", "FORE"):
        if side == "AFT":
            me = me_aft
        else:
            me = me_aft.copy()
            me.name = "DoorArm | rail kit FORE"
            me.transform(_mx("FORE"))
            me.flip_normals()
        ob = bpy.data.objects.new("DoorArm_%s | rail kit" % side, me)
        c_rail.objects.link(ob)
        tray = bpy.data.objects.new("DoorArm_%s | chain tray" % side, None)
        tray.empty_display_size = 0.2
        c_rail.objects.link(tray)
        c_frame1 = d["carriage"]["travel_ends"][side][1 if side == "AFT" else 0]
        for i, (x, y, rz) in enumerate(chain_positions(d, side, c_frame1)):
            lo = bpy.data.objects.new("DoorArm_%s | chain link %02d" % (side, i), chain_me)
            lo.parent = tray
            lo.location = (x, y, (d["energy_chain"][side]["z_bottom"] + d["energy_chain"][side]["z_top"]) / 2)
            lo.rotation_euler = (0, 0, rz)
            c_rail.objects.link(lo)
    doors = {"AFT": bpy.data.objects.get("Moonpool_Door_Aft"), "FORE": bpy.data.objects.get("Moonpool_Door_Fore")}
    for side, door in doors.items():
        if door is None:
            continue
        for ver, (L, att) in fixture_links(d, door, side).items():
            me = mesh_from_link(L, "DoorArm_%s | door fixture %s" % (side, ver))
            ob = bpy.data.objects.new("DoorArm_%s | door fixture %s" % (side, ver), me)
            ob.location = att
            ob["version"] = ver
            ob["attach_local"] = d["doors"][side]["attach_local"]
            c_fix.objects.link(ob)