"""Door robot arm: Cycles pictures (Coder stage 3; stage 'render' of dra_build.py) into previews/cycles/.

Nothing is saved: the file is opened, posed by its own drivers and rendered.
Views (--views, comma separated; default all):
  show       the showroom at the angles of ref1, ref2 and ref3: both versions side by side, and each alone
             (cycles_ref1_both.png, cycles_ref1_A.png, cycles_ref1_B.png, ...)
  close      close-ups like ref4 (turret and carriage), ref5 (lower arm), ref6 (wrist, A and B), ref7 (from behind,
             the lower arm stood upright as ref7 shows it)
  ship       the ship-fit scene, frames 1, 12, 30, 48, 90, both versions, both doors (the light rig saved in the file);
             narrow it with door:AFT, ver:A, frame:30
  samples:N, res:W   quality (default 128 samples)
Then contact sheets: cycles_ship_<door>_switch.png (A row over B row).
The showroom is lit by its own studio lights and a Poly Haven studio HDRI (studio_small_09, CC0, source_polyhaven/).
"""
import math
import os
import subprocess
import bpy
from mathutils import Vector, Matrix

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
MAGICK = r"E:\Tools\ImageMagick-7.1.2-Q16-HDRI\magick.exe"
HDRI = os.path.join(ROOT, "source_polyhaven", "hdris", "studio_small_09_2k.hdr")
FRAMES = (1, 12, 30, 48, 90)
MARKS = {1: "closed", 12: "retracted inboard", 30: "sliding open", 48: "open, in the pocket", 90: "sliding shut"}
CAMS = "90 STUDIO | cameras (not part of the asset)"


def gpu(sc, samples):
    sc.render.engine = 'CYCLES'
    try:
        prefs = bpy.context.preferences.addons["cycles"].preferences
        for kind in ('OPTIX', 'CUDA', 'HIP', 'ONEAPI'):
            try:
                prefs.compute_device_type = kind
            except Exception:
                continue
            (prefs.refresh_devices if hasattr(prefs, "refresh_devices") else prefs.get_devices)()
            if any(d.type == kind for d in prefs.devices):
                for d in prefs.devices:
                    d.use = d.type == kind
                sc.cycles.device = 'GPU'
                break
    except Exception as e:
        print("[dra] GPU not set:", e)
    sc.cycles.samples = samples
    sc.cycles.use_denoising = True
    sc.cycles.max_bounces = 8
    sc.cycles.glossy_bounces = 4
    sc.render.image_settings.file_format = 'PNG'
    sc.render.film_transparent = False
    sc.view_settings.view_transform = 'AgX'
    sc.view_settings.look = 'None'
    sc.view_settings.exposure = 0.0


def cam_coll(sc):
    c = bpy.data.collections.get(CAMS) or bpy.data.collections.new(CAMS)
    if c.name not in [x.name for x in sc.collection.children]:
        sc.collection.children.link(c)
    return c


def camera(sc, name):
    cd = bpy.data.cameras.get(name) or bpy.data.cameras.new(name)
    cd.clip_start = 0.05
    cd.clip_end = 100.0
    cam = bpy.data.objects.get(name) or bpy.data.objects.new(name, cd)
    c = cam_coll(sc)
    if cam.name not in c.objects:
        c.objects.link(cam)
    sc.camera = cam
    return cam


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


def fit_along(sc, cam, pts, direction, res, margin=1.08):
    """the camera looks along -direction at the middle of pts, at the distance that fits them all (perspective), or
    with the ortho scale that does (orthographic)"""
    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
    d = Vector(direction).normalized()
    cam.location = ctr + d * 20.0
    look(cam, ctr)
    bpy.context.view_layer.update()
    M = cam.matrix_world.to_3x3()
    r, u, f = M.col[0], M.col[1], -M.col[2]
    asp = res[1] / res[0]
    if cam.data.type == 'ORTHO':
        w = max(max(abs((p - ctr).dot(r)) for p in pts), max(abs((p - ctr).dot(u)) for p in pts) / asp)
        cam.data.ortho_scale = 2 * w * margin
        return cam
    tan_h = (cam.data.sensor_width / 2) / cam.data.lens
    tan_v = tan_h * asp
    dist = 1.0
    for p in pts:                              # the distance at which this point sits on the frame's edge
        v = p - ctr
        x, y, z = v.dot(r), v.dot(u), v.dot(f)
        dist = max(dist, abs(x) / tan_h - z, abs(y) / tan_v - z)
    cam.location = ctr - f * dist * margin
    return cam


def arm_points(sc, v):
    dg = bpy.context.evaluated_depsgraph_get()
    pts = []
    for o in sc.objects:
        if o.name.startswith("Showroom | arm %s |" % v) and o.type == 'MESH':
            ev = o.evaluated_get(dg)
            pts += [ev.matrix_world @ Vector(c) for c in ev.bound_box]
    return pts


def studio(sc):
    """the showroom's world: kept in the file by the surf stage (dra_surf.studio_look); this only makes sure there is
    one (an HDRI under the scene's own lights) when the file has none"""
    w = bpy.data.worlds.get("Showroom | world")
    if w is not None and w.use_nodes and any(n.type == 'TEX_ENVIRONMENT' for n in w.node_tree.nodes):
        sc.world = w
        return
    w = w or bpy.data.worlds.new("Showroom | world")
    w.use_nodes = True
    nt = w.node_tree
    bg = nt.nodes.get("Background")
    if os.path.exists(HDRI) and bg is not None:
        env = next((n for n in nt.nodes if n.type == 'TEX_ENVIRONMENT'), None) or nt.nodes.new('ShaderNodeTexEnvironment')
        env.image = bpy.data.images.load(HDRI, check_existing=True)
        nt.links.new(env.outputs["Color"], bg.inputs["Color"])
        bg.inputs["Strength"].default_value = 0.30
    sc.world = w


def shot(sc, path, res):
    sc.render.resolution_x, sc.render.resolution_y = res
    sc.render.resolution_percentage = 100
    sc.render.filepath = path
    bpy.ops.render.render(write_still=True, scene=sc.name)
    print("[dra] rendered", path, flush=True)


def showroom(views, outdir, samples):
    import dra_preview as P
    sc = bpy.data.scenes["DoorArm | Showroom"]
    wm = bpy.context.window_manager
    if wm.windows:
        wm.windows[0].scene = sc
    gpu(sc, samples)
    studio(sc)
    done = []
    tur = bpy.data.objects.get("Showroom | arm A | turret")
    low = bpy.data.objects.get("Showroom | arm A | lower arm")
    cyc = bpy.data.objects.get("Showroom | cyclorama")
    if "show" in views:
        for ref in ("ref1", "ref2", "ref3"):
            for which in ("A", "B", "both"):
                P.show_only(None if which == "both" else which)
                cam = P.camera_for(ref, sc)               # the reference's angle (dra_preview)
                if which != "A":                           # framed on what is shown, along the same direction
                    d = (cam.matrix_world.translation - (cam.matrix_world.translation + cam.matrix_world.to_3x3().col[2] * -1))
                    direction = cam.matrix_world.to_3x3().col[2]       # from the target toward the camera
                    pts = arm_points(sc, "A") + arm_points(sc, "B") if which == "both" else arm_points(sc, "B")
                    fit_along(sc, cam, pts, direction, (1920, 1080) if which == "both" else (1400, 1400))
                sc.camera = cam
                if cyc:
                    cyc.hide_render = cam.location.y < -2.5
                res = (1920, 1080) if which == "both" else (1400, 1400)
                p = os.path.join(outdir, "cycles_%s_%s.png" % (ref, which))
                shot(sc, p, res)
                done.append(p)
    if "close" in views:
        for ref, ver in (("ref4", "A"), ("ref5", "A"), ("ref6A", "A"), ("ref6B", "B"), ("ref7", "A")):
            P.show_only(ver)
            if tur:
                tur.rotation_euler = (0, 0, math.radians(90) if ref == "ref7" else 0.0)
            if low:
                low.rotation_euler = (0, math.radians(P.REF7_J2) if ref == "ref7" else 0.0, 0)
            cam = P.camera_for(ref, sc)
            sc.camera = cam
            if cyc:
                cyc.hide_render = cam.location.y < -2.5
            p = os.path.join(outdir, "cycles_close_%s.png" % ref)
            shot(sc, p, (1400, 1400))
            done.append(p)
        if tur:
            tur.rotation_euler = (0, 0, 0)
        if low:
            low.rotation_euler = (0, 0, 0)
    P.show_only(None)
    if cyc:
        cyc.hide_render = False
    return done


def ship(views, outdir, samples):
    import dra_render as RR
    sf = bpy.data.scenes["DoorArm | Ship fit"]
    wm = bpy.context.window_manager
    if wm.windows:
        wm.windows[0].scene = sf
    gpu(sf, samples)
    ctl = bpy.data.objects["DoorArm | CONTROLS"]
    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)
    res = (1600, 900)
    done = []

    def setv(v):
        ctl["version"] = 0 if v == "A" else 1
        ctl.update_tag()
    for key in doors:
        for f in frames:
            pts = []
            for v in ("A", "B"):
                setv(v)
                sf.frame_set(f)
                pts += RR.corners(sf, [o for o in sf.objects if o.name.startswith("DoorArm_%s_%s |" % (key, v))
                                       and o.type == 'MESH'])
                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)]
            cam = camera(sf, "Cam | %s pocket" % key)
            side = 1.0 if key == "AFT" else -1.0
            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
            cam.location = Vector((ctr.x + side * 1.6, RR.CAM_Y, RR.CAM_Z))
            cam.data.lens = 20.0
            look(cam, ctr)
            bpy.context.view_layer.update()
            M = cam.matrix_world.to_3x3()
            r, u, fwd = M.col[0], M.col[1], -M.col[2]
            mxx = myy = 1e-6
            for p in pts:
                vv = p - cam.location
                dd = vv.dot(fwd)
                if dd > 0.05:
                    mxx = max(mxx, abs(vv.dot(r)) / dd)
                    myy = max(myy, abs(vv.dot(u)) / dd)
            sw = cam.data.sensor_width
            cam.data.sensor_fit = 'HORIZONTAL'
            cam.data.lens = min(sw / 2 / mxx, sw * res[1] / res[0] / 2 / myy) / 1.08
            for v in vers:
                setv(v)
                sf.frame_set(f)
                p = os.path.join(outdir, "cycles_ship_%s_%s_f%03d.png" % (key, v, f))
                shot(sf, p, res)
                done.append(p)
    setv("A")
    sf.frame_set(1)
    return done


def sheets(outdir):
    if not os.path.exists(MAGICK):
        return
    for key in ("AFT", "FORE"):
        rows = []
        for v in ("A", "B"):
            files = [os.path.join(outdir, "cycles_ship_%s_%s_f%03d.png" % (key, v, f)) for f in FRAMES]
            if not all(os.path.exists(p) for p in files):
                continue
            tiles = []
            for f, p in zip(FRAMES, files):
                t = os.path.join(outdir, "_tile_%s_%s_%03d.png" % (key, v, f))
                subprocess.run([MAGICK, p, "-resize", "640x360", "-gravity", "NorthWest", "-fill", "white", "-undercolor",
                                "#00000090", "-pointsize", "17", "-annotate", "+8+8",
                                "%s  version %s  frame %d  (%s)" % (key, v, f, MARKS[f]), t], check=False, timeout=120)
                tiles.append(t)
            row = os.path.join(outdir, "_row_%s_%s.png" % (key, v))
            subprocess.run([MAGICK] + tiles + ["+append", row], check=False, timeout=120)
            rows.append(row)
            for t in tiles:
                if os.path.exists(t):
                    os.remove(t)
        if len(rows) == 2:
            out = os.path.join(outdir, "cycles_ship_%s_switch.png" % key)
            subprocess.run([MAGICK] + rows + ["-append", out], check=False, timeout=120)
            print("[dra] sheet", out, flush=True)
        for r in rows:
            if os.path.exists(r):
                os.remove(r)


def render(views, outdir, res=None):
    outdir = os.path.join(outdir, "cycles") if os.path.basename(outdir) != "cycles" else outdir
    os.makedirs(outdir, exist_ok=True)
    views = views or ["show", "close", "ship"]
    smp = [int(v.split(":")[1]) for v in views if v.startswith("samples:")]
    samples = smp[0] if smp else 128
    done = []
    if "show" in views or "close" in views:
        done += showroom(views, outdir, samples)
    if "ship" in views:
        done += ship(views, outdir, samples)
    sheets(outdir)
    return done