"""Door robot arm: Workbench preview pictures of the showroom (shape checks against the references).

ref1 is rendered orthographic at the reference's own scale (0.0022 m per ref1 pixel, 2000 px = 4.4 m) with the arm frame
origin where ref1 has it (x 1000, y 1550), so the picture can be laid over ref1 pixel for pixel (an overlay is written
next to it when ImageMagick is there). The other views are perspective cameras set to look like the references.
"""
import math
import os
import subprocess
import bpy
from mathutils import Vector, Matrix

ROOT_Z = 1.012
REF7_J2 = 33.6                  # lower arm rotation (deg, about +Y) that stands it upright for the ref7 picture
MAGICK = r"E:\Tools\ImageMagick-7.1.2-Q16-HDRI\magick.exe"
REFDIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "reference")


def arm_x(v):
    ob = bpy.data.objects.get("Showroom | arm " + v)
    return ob.matrix_world.translation.x if ob else (-1.9 if v == "A" else 1.9)


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


def setup(sc, res=1000):
    sc.render.engine = 'BLENDER_WORKBENCH'
    sh = sc.display.shading
    sh.light = 'STUDIO'
    sh.color_type = 'MATERIAL'
    sh.show_cavity = True
    sh.cavity_type = 'BOTH'
    sh.cavity_ridge_factor = 1.0
    sh.cavity_valley_factor = 1.0
    sh.curvature_ridge_factor = 1.0
    sh.curvature_valley_factor = 1.0
    sh.show_object_outline = True
    sh.object_outline_color = (0.05, 0.05, 0.05)
    sh.show_specular_highlight = True
    sc.render.resolution_x = res
    sc.render.resolution_y = res
    sc.render.resolution_percentage = 100
    sc.render.film_transparent = False
    sc.display.render_aa = '8'
    sc.render.image_settings.file_format = 'PNG'


def show_only(version):
    """hide the other arm (A or B); None shows both"""
    for ob in bpy.data.objects:
        v = ob.get("version")
        if ob.name.startswith("Showroom | arm") and v:
            hide = version is not None and v != version
            ob.hide_render = hide
            ob.hide_viewport = hide


VIEWS = {
    # name: (version shown, camera function)
    "ref1": "A", "ref2": "A", "ref3": "A", "ref4": "A", "ref5": "A", "ref6A": "A", "ref6B": "B", "ref7": "A", "both": None,
    "ref6A_side": "A", "ref6B_side": "B",
}


def camera_for(name, sc):
    cam_data = bpy.data.cameras.get("Preview cam") or bpy.data.cameras.new("Preview cam")
    cam = bpy.data.objects.get("Preview cam") or bpy.data.objects.new("Preview cam", cam_data)
    if cam.name not in sc.collection.objects:
        sc.collection.objects.link(cam)
    sc.camera = cam
    cd = cam.data
    cd.type = 'PERSP'
    cd.clip_end = 100.0
    import dra_rigdata as _rd
    SCL = _rd.S                                  # scale-up: the views are set in build units and scaled at the end
    xa = arm_x("A") / SCL
    xb = arm_x("B") / SCL
    o = Vector((xa, 0, ROOT_Z))                  # arm A frame origin
    if name == "ref1":
        cd.type = 'ORTHO'
        cd.ortho_scale = 4.4
        cam.location = o + Vector((0.0, 12.0, 1.21))
        cam.rotation_euler = (math.radians(90), 0, math.radians(180))
    elif name == "ref2":
        cd.lens = 85
        tgt = o + Vector((-0.05, 0, 1.30))
        cam.location = tgt + Vector((7.1, 9.1, 0.25))
        look(cam, tgt)
    elif name == "ref3":
        cd.lens = 85
        tgt = o + Vector((-0.20, 0, 1.30))
        cam.location = tgt + Vector((-5.6, -8.6, 0.25))
        look(cam, tgt)
    elif name == "ref4":
        cd.lens = 85
        tgt = o + Vector((0.10, 0.0, 0.25))
        cam.location = tgt + Vector((2.3, 4.0, 0.50))
        look(cam, tgt)
    elif name == "ref5":
        # set like ref5 (the critique's matched camera): from the front-left and a little below, so the lower arm's
        # dark core shows its front face and the knee cap reads as an ellipse
        cd.lens = 85
        tgt = o + Vector((-1.0, 0.2, 1.30))
        az, el = math.radians(62), math.radians(-12)
        cam.location = tgt + 4.4 * Vector((math.cos(el) * math.cos(az), math.cos(el) * math.sin(az), math.sin(el)))
        look(cam, tgt)
    elif name in ("ref6A", "ref6B"):
        # set like ref6's photo (CRITIQUE_1 item 1: the hub's face reads as a narrow ellipse, the tool axis about 25 deg
        # out of the picture): 3.8 m from the forearm, azimuth 62 deg from the tool axis toward the arm's left, 12 deg below
        cd.lens = 85
        base = Vector((xa if name == "ref6A" else xb, 0, ROOT_Z))
        tgt = base + Vector((0.80, 0.0, 2.06))
        az, el = math.radians(62), math.radians(-12)
        cam.location = tgt + 3.8 * Vector((math.cos(el) * math.cos(az), math.cos(el) * math.sin(az), math.sin(el)))
        look(cam, tgt)
    elif name in ("ref6A_side", "ref6B_side"):
        cd.lens = 85
        base = Vector((xa if name == "ref6A_side" else xb, 0, ROOT_Z))
        tgt = base + Vector((0.8, 0.0, 2.05))
        cam.location = tgt + Vector((0.0, 4.2, 0.0))
        look(cam, tgt)
    elif name == "ref7":
        cd.lens = 85
        tgt = o + Vector((0, 0, 0.45))
        cam.location = tgt + Vector((0.0, -5.2, 0.15))
        look(cam, tgt)
    elif name == "both":
        cd.lens = 50
        tgt = Vector(((xa + xb) / 2, 0, 1.9))
        cam.location = tgt + Vector((2.0, 11.0, 1.0))
        look(cam, tgt)
    cam.location = cam.location * SCL
    if cd.type == 'ORTHO':
        cd.ortho_scale *= SCL
    cd.clip_end = 100.0 * SCL
    return cam


def render(views, outdir, res=1000):
    os.makedirs(outdir, exist_ok=True)
    sc = bpy.data.scenes["DoorArm | Showroom"]
    wm = bpy.context.window_manager
    if wm.windows:
        wm.windows[0].scene = sc               # the rigged file opens in the ship scene
    setup(sc, res)
    views = views or ["ref1", "ref2", "ref3", "ref4", "ref5", "ref6A", "ref6B", "ref7", "both", "shipAFT", "shipAFT_rail"]
    tur = bpy.data.objects.get("Showroom | arm A | turret")
    low = bpy.data.objects.get("Showroom | arm A | lower arm")
    for v in views:
        if v.startswith("ship") or v not in VIEWS:
            continue
        show_only(VIEWS.get(v, "A"))
        if tur:
            tur.rotation_euler = (0, 0, math.radians(90) if v == "ref7" else 0.0)
        if low:
            # ref7 shows the lower arm standing upright over the drum (the rest pose, ref1's, leans it 33.6 deg back,
            # which would hide the drum from behind): J2 turned for this picture only
            low.rotation_euler = (0, math.radians(REF7_J2) if v == "ref7" else 0.0, 0)
        cam = camera_for(v, sc)
        cyc = bpy.data.objects.get("Showroom | cyclorama")
        if cyc:
            cyc.hide_render = cam.location.y < -2.5
        sc.render.filepath = os.path.join(outdir, "geo_%s.png" % v)
        bpy.ops.render.render(write_still=True, scene=sc.name)
        print("[dra] rendered", sc.render.filepath, flush=True)
        if v == "ref1" and os.path.exists(MAGICK):
            ref = os.path.join(REFDIR, "ref1_side_left.jpg")
            out = os.path.join(outdir, "geo_ref1_overlay.png")
            try:
                subprocess.run([MAGICK, ref, "-resize", "%dx%d" % (res, res), sc.render.filepath, "-compose", "blend",
                                "-define", "compose:args=50,50", "-composite", out], check=True, timeout=60)
            except Exception as e:
                print("[dra] overlay failed", e)
    if tur:
        tur.rotation_euler = (0, 0, 0)
    if low:
        low.rotation_euler = (0, 0, 0)
    show_only(None)
    for v in views:
        if v.startswith("ship"):
            ship_view(v, outdir)
    if "sheet" in views or all(os.path.exists(os.path.join(outdir, "geo_%s.png" % v)) for v, _r, _l in SHEET):
        compare_sheet(outdir)


# the comparison sheet: each reference beside the model's view of it, two pairs per row
SHEET = [("ref1", "ref1_side_left.jpg", "ref1"), ("ref2", "ref2_three_quarter_left.jpg", "ref2"),
         ("ref3", "ref3_three_quarter_right.jpg", "ref3"), ("ref4", "ref4_base_wire_closeup.jpg", "ref4"),
         ("ref5", "ref5_lower_arm_wire_closeup.jpg", "ref5"), ("ref7", "ref7_front_base_wire.jpg", "ref7 (lower arm upright)"),
         ("ref6A", "ref6_wrist_redline_no_grab_limbs.png", "ref6 / version A"),
         ("ref6B", "ref6_wrist_redline_no_grab_limbs.png", "ref6 red sketch / version B")]


def compare_sheet(outdir):
    if not os.path.exists(MAGICK):
        return
    tmp = os.path.join(outdir, "_sheet_tmp")
    os.makedirs(tmp, exist_ok=True)
    pairs = []
    for i, (v, ref, lab) in enumerate(SHEET):
        a, b, pr = (os.path.join(tmp, "%s_%d.png" % (x, i)) for x in ("a", "b", "p"))
        for src, dst, text in ((os.path.join(REFDIR, ref), a, "reference: " + lab),
                               (os.path.join(outdir, "geo_%s.png" % v), b, "model: geo_%s.png" % v)):
            subprocess.run([MAGICK, src, "-resize", "500x500", "-gravity", "NorthWest", "-fill", "white", "-undercolor",
                            "#00000099", "-pointsize", "18", "-annotate", "+6+6", text, "+repage", dst], check=False,
                           timeout=120)
        subprocess.run([MAGICK, a, b, "+append", "+repage", pr], check=False, timeout=120)
        pairs.append(pr)
    rows = []
    for k in range(0, len(pairs), 2):
        r = os.path.join(tmp, "row_%d.png" % k)
        subprocess.run([MAGICK, pairs[k], pairs[k + 1], "+append", "+repage", r], check=False, timeout=120)
        rows.append(r)
    out = os.path.join(outdir, "geo_compare_sheet.png")
    subprocess.run([MAGICK] + rows + ["-append", "+repage", out], check=False, timeout=120)
    for f in os.listdir(tmp):
        os.remove(os.path.join(tmp, f))
    os.rmdir(tmp)
    print("[dra] sheet", out, flush=True)


def ship_view(v, outdir):
    """the ship scene at frame 1 (no arms yet): the rail kit, chain and door fixtures in the pocket"""
    sf = bpy.data.scenes["DoorArm | Ship fit"]
    setup(sf)
    sf.frame_set(1)
    cam_data = bpy.data.cameras.get("Preview cam ship") or bpy.data.cameras.new("Preview cam ship")
    cam = bpy.data.objects.get("Preview cam ship") or bpy.data.objects.new("Preview cam ship", cam_data)
    if cam.name not in sf.collection.objects:
        sf.collection.objects.link(cam)
    sf.camera = cam
    cam_data.lens = 24
    cam_data.clip_start = 0.05
    side = "AFT" if "AFT" in v else "FORE"
    x = -5.6 if side == "AFT" else 11.7
    if v.endswith("_rail"):
        cam.location = (x + (2.6 if side == "AFT" else -2.6), -4.3, -4.45)
        look(cam, (x - (1.0 if side == "AFT" else -1.0), -5.6, -3.95))
    else:
        cam.location = (x, -1.2, -5.2)
        look(cam, (x + (2.0 if side == "AFT" else -2.0), -5.9, -4.6))
    for o in sf.objects:
        if o.name.startswith("Moonpool | deck") or o.name.startswith("Moonpool | hull lining"):
            o.hide_render = False
    sf.render.filepath = os.path.join(outdir, "geo_%s.png" % v)
    bpy.ops.render.render(write_still=True, scene=sf.name)
    print("[dra] rendered", sf.render.filepath, flush=True)