"""Door robot arm: Cycles pictures of the rigged arms in the ship scene (stage 'rig_render' -> previews/rig/).

The file is opened, posed by its own drivers and rendered; nothing is saved. What this stage adds lives only in the
session: a light rig in the collection 'DoorArm | render lights (not exported)' (the pocket under the deck is dark:
two long area lights under the deck inboard of each rail and a soft fill from the moon pool well; the docking lamps
keep their own emission), cameras, and plain stand-in colours for the ship materials whose image files are missing
in the ship (an old absolute path there; not ours to fix).

Views (--views, comma separated; default all):
  shots     both doors x both versions x frames 1, 12, 30, 48, 90: one camera per frame (the same for A and B),
            inboard of the pocket under the deck, framing both versions' arms and the door round the hold
  grip      version A's gripper closed on the door's knob (frame 1, AFT), close
  grip_wb   with grip: the same picture in Workbench into previews/geo/geo_gripA_closed_on_knob.png
  plate     version B's plate on its pad (frame 1, AFT), close
  door:AFT, ver:A, frame:30, samples:16   narrow a run down (quick tests)
Then contact sheets: rig_<door>_<version>_strip.png (the five frames), rig_switch_<door>.png (A over B per frame).
"""
import math
import os
import subprocess
import bpy
from mathutils import Vector, Matrix

MAGICK = r"E:\Tools\ImageMagick-7.1.2-Q16-HDRI\magick.exe"
FRAMES = (1, 12, 30, 48, 90)
MARKS = {1: "closed", 12: "retracted inboard", 30: "sliding open", 48: "open, in the pocket", 90: "sliding shut"}
LIGHTS = "DoorArm | render lights (not exported)"


def look(ob, target):
    d = (Vector(target) - ob.location).normalized()
    ob.rotation_euler = d.to_track_quat('-Z', 'Y').to_euler()


def setup(sc, res=(1280, 720), samples=96):
    sc.render.engine = 'CYCLES'
    try:
        prefs = bpy.context.preferences.addons["cycles"].preferences
        prefs.compute_device_type = 'HIP'
        prefs.get_devices()
        for d in prefs.devices:
            d.use = True
        sc.cycles.device = 'GPU'
    except Exception as e:
        print("[dra] GPU not set:", e)
    sc.cycles.samples = samples
    sc.cycles.use_denoising = True
    sc.cycles.max_bounces = 6
    sc.render.resolution_x, sc.render.resolution_y = res
    sc.render.resolution_percentage = 100
    sc.render.image_settings.file_format = 'PNG'
    sc.render.film_transparent = False
    sc.view_settings.view_transform = 'AgX'
    sc.view_settings.look = 'None'
    w = bpy.data.worlds.get("DoorArm | pocket world") or bpy.data.worlds.new("DoorArm | pocket world")
    w.use_nodes = True
    bg = w.node_tree.nodes.get("Background")
    if bg:
        bg.inputs[0].default_value = (0.018, 0.021, 0.026, 1)
        bg.inputs[1].default_value = 1.0
    sc.world = w


def stand_in_materials():
    """ship materials whose images are missing render magenta: give them plain colours for these pictures"""
    looks = {"PBR | Brushed titanium": ((0.50, 0.51, 0.53, 1), 1.0, 0.32),
             "PBR | Graphite lower hull": ((0.06, 0.065, 0.07, 1), 0.6, 0.45)}
    for m in bpy.data.materials:
        if not m.node_tree:
            continue
        missing = [n for n in m.node_tree.nodes if n.type == 'TEX_IMAGE' and n.image is not None and not n.image.has_data
                   and not os.path.exists(bpy.path.abspath(n.image.filepath))]
        if not missing:
            continue
        bsdf = next((n for n in m.node_tree.nodes if n.type == 'BSDF_PRINCIPLED'), None)
        for n in missing:
            for out in n.outputs:
                for lk in list(out.links):
                    m.node_tree.links.remove(lk)
        if bsdf is not None:
            col, met, rough = looks.get(m.name, ((0.35, 0.36, 0.38, 1), 0.5, 0.4))
            bsdf.inputs["Base Color"].default_value = col
            bsdf.inputs["Metallic"].default_value = met
            bsdf.inputs["Roughness"].default_value = rough
        print("[dra] stand-in colours for", m.name)


def light_rig(sf):
    c = bpy.data.collections.get(LIGHTS) or bpy.data.collections.new(LIGHTS)
    if c.name not in [x.name for x in sf.collection.children]:
        sf.collection.children.link(c)
    c["not_exported"] = True

    def area(name, loc, size, energy, target, color=(1.0, 0.97, 0.93)):
        ld = bpy.data.lights.get(name) or bpy.data.lights.new(name, 'AREA')
        ld.shape = 'RECTANGLE'
        ld.size, ld.size_y = size
        ld.energy = energy
        ld.color = color
        o = bpy.data.objects.get(name) or bpy.data.objects.new(name, ld)
        if o.name not in c.objects:
            c.objects.link(o)
        o.location = loc
        look(o, target)
        return o
    for key, sx in (("AFT", 1.0), ("FORE", -1.0)):
        x0 = mx(key, -5.3)                     # the pocket's whole length, end wall to the rail's inner end
        area("Light | %s under deck inboard" % key, (x0, -4.55, -3.90), (7.5, 1.0), 750, (x0, -6.6, -6.2))
        area("Light | %s under deck outboard" % key, (x0, -8.60, -3.95), (7.5, 0.6), 300, (x0, -6.0, -6.0),
             color=(0.90, 0.95, 1.0))
        area("Light | %s fill from the well" % key, (mx(key, -2.1), -2.6, -5.6), (3.0, 3.0), 380, (mx(key, -5.0), -7.0, -6.0),
             color=(0.85, 0.92, 1.0))
        area("Light | %s fill over the moon pool" % key, (mx(key, 0.5), -1.5, -4.6), (3.0, 2.0), 280,
             (mx(key, -1.5), -7.0, -6.5), color=(0.85, 0.92, 1.0))
    return c


def camera(sf, name, loc, target, lens):
    cd = bpy.data.cameras.get(name) or bpy.data.cameras.new(name)
    cd.lens = lens
    cd.clip_start = 0.05
    cd.clip_end = 60.0
    cam = bpy.data.objects.get(name) or bpy.data.objects.new(name, cd)
    if cam.name not in bpy.data.collections[LIGHTS].objects:
        bpy.data.collections[LIGHTS].objects.link(cam)
    cam.location = loc
    look(cam, target)
    sf.camera = cam
    return cam


def mx(key, x):
    return x if key == "AFT" else 6.1 - x


DOOR = {"AFT": "Moonpool_Door_Aft", "FORE": "Moonpool_Door_Fore"}
CAM_Y, CAM_Z = -1.6, -4.35            # in the pocket's open space above the door, inboard, just under the deck


def corners(sf, obs):
    dg = bpy.context.evaluated_depsgraph_get()
    out = []
    for o in obs:
        ev = o.evaluated_get(dg)
        out += [ev.matrix_world @ Vector(c) for c in ev.bound_box]
    return out


def fit_camera(sf, name, key, pts, res):
    """a camera inboard of the pocket (under the deck, in the basin's open space), a little toward the moon pool, aimed
    at the middle of pts, its lens set so every point is in the picture with a margin"""
    lo = Vector((min(p.x for p in pts), min(p.y for p in pts), min(p.z for p in pts)))
    hi = Vector((max(p.x for p in pts), max(p.y for p in pts), max(p.z for p in pts)))
    ctr = (lo + hi) / 2
    side = 1.0 if key == "AFT" else -1.0
    loc = Vector((ctr.x + side * 1.6, CAM_Y, CAM_Z))
    cam = camera(sf, name, loc, ctr, 20.0)
    bpy.context.view_layer.update()
    M = cam.matrix_world.to_3x3()
    r, u, f = M.col[0], M.col[1], -M.col[2]
    mxx = myy = 1e-6
    for p in pts:
        v = p - loc
        d = v.dot(f)
        if d <= 0.05:
            continue
        mxx = max(mxx, abs(v.dot(r)) / d)
        myy = max(myy, abs(v.dot(u)) / d)
    sw = cam.data.sensor_width
    sh = sw * res[1] / res[0]
    cam.data.sensor_fit = 'HORIZONTAL'
    cam.data.lens = min(sw / 2 / mxx, sh / 2 / myy) / 1.08
    return cam


def set_version(v):
    ctl = bpy.data.objects["DoorArm | CONTROLS"]
    ctl["version"] = 0 if v == "A" else 1
    ctl.update_tag()


def shot(sf, path):
    sf.render.filepath = path
    bpy.ops.render.render(write_still=True, scene=sf.name)
    print("[dra] rendered", path, flush=True)


def render(sf, views, outdir, res=None):
    os.makedirs(outdir, exist_ok=True)
    views = views or ["shots", "grip", "plate"]
    smp = [int(v.split(":")[1]) for v in views if v.startswith("samples:")]
    setup(sf, (res, int(res * 9 / 16)) if res else (1280, 720), samples=smp[0] if smp else 96)
    stand_in_materials()
    light_rig(sf)
    done = []
    doors = [v.split(":")[1] for v in views if v.startswith("door:")] or ["AFT", "FORE"]
    vers = [v.split(":")[1] for v in views if v.startswith("ver:")] or ["A", "B"]
    frames = [int(v.split(":")[1]) for v in views if v.startswith("frame:")] or list(FRAMES)
    if "shots" in views:
        for key in doors:
            for f in frames:
                # one camera per frame for both versions, framing both versions' arms and the door at that frame
                pts = []
                for v in ("A", "B"):
                    set_version(v)
                    sf.frame_set(f)
                    pts += corners(sf, [o for o in sf.objects if o.name.startswith("DoorArm_%s_%s |" % (key, v))
                                        and o.type == 'MESH'])
                    # and the door round the place the arm holds (0.9 m every way), so the door reads as the door
                    fx = bpy.data.objects["DoorArm_%s | door fixture %s" % (key, v)]
                    at = fx.evaluated_get(bpy.context.evaluated_depsgraph_get()).matrix_world.translation
                    pts += [at + Vector((dx, dy, dz)) for dx in (-0.9, 0.9) for dy in (-0.9, 0.9) for dz in (-0.9, 0.9)]
                fit_camera(sf, "Cam | %s inboard" % key, key, pts, (res, int(res * 9 / 16)) if res else (1280, 720))
                for v in vers:
                    set_version(v)
                    sf.frame_set(f)
                    p = os.path.join(outdir, "rig_%s_%s_f%03d.png" % (key, v, f))
                    shot(sf, p)
                    done.append(p)
    if "grip" in views:
        set_version("A")
        sf.frame_set(1)
        dg = bpy.context.evaluated_depsgraph_get()
        fl = bpy.data.objects["DoorArm_AFT_A | flange"].evaluated_get(dg).matrix_world
        wc = fl.translation
        ax = fl.to_3x3().col[0]
        tgt = wc + ax * 0.37
        # from the moon pool side, across the tool axis and a little off the door: the hub, the four fingers closed
        # round the knob's head, the knob's stem and pad on the door in profile
        camera(sf, "Cam | AFT grip", tgt + Vector((0.95, 0.0, 0.0)) - ax * 0.30, tgt, 50)
        shot(sf, os.path.join(outdir, "rig_gripA_closed_on_knob.png"))
        done.append("grip")
    if "plate" in views:
        set_version("B")
        sf.frame_set(1)
        dg = bpy.context.evaluated_depsgraph_get()
        fl = bpy.data.objects["DoorArm_AFT_B | flange"].evaluated_get(dg).matrix_world
        wc = fl.translation
        ax = fl.to_3x3().col[0]
        tgt = wc + ax * 0.22
        camera(sf, "Cam | AFT plate", tgt + Vector((0.75, 1.1, 0.6)), tgt, 40)
        shot(sf, os.path.join(outdir, "rig_plateB_on_pad.png"))
    if "grip" in views and "grip_wb" in views:
        # CRITIQUE_1 item 11: the closed grip on the knob drawn from the real meshes (Workbench, like the other geo
        # previews; it replaces the sweep's flat-shaded picture in previews/geo)
        import dra_preview
        set_version("A")
        sf.frame_set(1)
        sf.camera = bpy.data.objects["Cam | AFT grip"]
        dra_preview.setup(sf, 1000)
        sf.render.resolution_x, sf.render.resolution_y = 1280, 720
        geo_dir = os.path.join(os.path.dirname(outdir), "geo")
        shot(sf, os.path.join(geo_dir, "geo_gripA_closed_on_knob.png"))
    set_version("A")
    sf.frame_set(1)
    sheets(outdir)
    return done


def sheets(outdir):
    """strips of the five frames per door and version, and the switch (A above B) per door"""
    try:
        from PIL import Image, ImageDraw
    except ImportError:
        Image = None
    for key in ("AFT", "FORE"):
        rows = {}
        for v in ("A", "B"):
            files = [os.path.join(outdir, "rig_%s_%s_f%03d.png" % (key, v, f)) for f in FRAMES]
            if all(os.path.exists(p) for p in files):
                rows[v] = files
                if os.path.exists(MAGICK):
                    out = os.path.join(outdir, "rig_%s_%s_strip.png" % (key, v))
                    tiles = []
                    for f, p in zip(FRAMES, files):
                        t = os.path.join(outdir, "_tile_%s_%s_%03d.png" % (key, v, f))
                        lab = "%s  version %s  frame %d  (%s)" % (key, v, f, MARKS.get(f, ""))
                        r = subprocess.run([MAGICK, p, "-resize", "640x360", "-gravity", "NorthWest", "-fill", "white",
                                            "-undercolor", "#00000090", "-pointsize", "17", "-annotate", "+8+8", lab, t],
                                           check=False, timeout=120)
                        tiles.append(t if r.returncode == 0 and os.path.exists(t) else p)
                    subprocess.run([MAGICK] + tiles + ["-resize", "640x360", "+append", out], check=False, timeout=120)
                    for t in tiles:
                        if os.path.basename(t).startswith("_tile_") and os.path.exists(t):
                            os.remove(t)
        if len(rows) == 2 and os.path.exists(MAGICK):
            out = os.path.join(outdir, "rig_switch_%s.png" % key)
            a = os.path.join(outdir, "rig_%s_A_strip.png" % key)
            b = os.path.join(outdir, "rig_%s_B_strip.png" % key)
            subprocess.run([MAGICK, a, b, "-append", out], check=False, timeout=120)
            print("[dra] sheet", out, flush=True)