"""Grouper: Cycles pictures (headless; nothing saved). The lights are on (Grouper_Root["lights"] = 1) in every picture.

blender.exe -b Tender_Grouper.blend --python build/gr_render.py -- TAG SAMPLES MODE [MODE ...]
  studio     bright studio light (a grey gradient world, key sun, fills; Standard view transform):
             previews/<TAG>_studio_<side|top|front|back|q_fore|q_aft|cabin_hatch|cabin_door>.png
  water      two underwater pictures (dark water, haze, the lamps as spot lights): previews/<TAG>_water_<nose|stern>.png
  concept    the concept render's fitted camera (build/checks/concept_cam.json), dark water light:
             previews/compare/concept_render.png (gr_pairs.py sets it beside the concept)
  lookalike  the sheet's views framed as the sheet frames them (orthographic, 2 render px per sheet px, studio light,
             transparent film): previews/compare/lk_<side|top|bottom|front|back>.png, and the cabin cut at the A plane
             (the clamp axis, Y -0.09, looking forward): lk_section.png
"""
import bpy
import json
import math
import os
import sys
from mathutils import Vector, Matrix, Euler

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
PREV = os.path.join(ROOT, "previews")
CMP = os.path.join(PREV, "compare")
os.makedirs(CMP, exist_ok=True)
args = sys.argv[sys.argv.index("--") + 1:] if "--" in sys.argv else []
TAG = args[0] if args else "v01"
SAMPLES = int(args[1]) if len(args) > 1 else 96
MODES = args[2:] or ["studio", "water", "concept", "lookalike"]
G = {}
exec(open(os.path.join(HERE, "gr_measure.py"), encoding="utf-8").read(), G)
S, KF, KB = G["S"], G["KF"], G["KB"]
D = bpy.data
sc = bpy.context.scene
root = D.objects["Grouper_Root"]


def rearm():
    for ob in D.objects:
        if ob.animation_data:
            for fc in ob.animation_data.drivers:
                fc.driver.expression = fc.driver.expression
                fc.driver.is_valid = True
            ob.update_tag()
    for m in D.materials:
        if m.node_tree and m.node_tree.animation_data:
            for fc in m.node_tree.animation_data.drivers:
                fc.driver.expression = fc.driver.expression
                fc.driver.is_valid = True
    sc.frame_set(sc.frame_current)
    bpy.context.view_layer.update()


def pose(**kw):
    for k, v in kw.items():
        root[k] = v
    root.update_tag()
    rearm()


pose(lights=1.0, hatch_open_deg=0.0, door_l_deg=0.0, door_r_deg=0.0, antenna_fold_deg=0.0, rotor_turns=0.0)
sc.render.engine = 'CYCLES'
sc.cycles.samples = SAMPLES
sc.cycles.use_denoising = True
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 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 ex:
    print("GPU not set", ex)
sc.view_settings.view_transform = 'Standard'
sc.view_settings.look = 'None'
sc.render.image_settings.file_format = 'PNG'
sc.render.resolution_percentage = 100
sc.cycles.max_bounces = 12
sc.cycles.transmission_bounces = 12
sc.cycles.transparent_max_bounces = 16
for o in D.objects:
    if o.name.startswith(("UCX_", "SOCKET_")):
        o.hide_render = True
coll = D.collections.new("_render")
sc.collection.children.link(coll)
cam = D.objects.new("_cam", D.cameras.new("_cam"))
coll.objects.link(cam)
sc.camera = cam
cam.data.clip_start = 0.02
cam.data.clip_end = 300.0
LIGHTS = []


def light(name, kind, energy, rot=(0, 0, 0), loc=(0, 0, 0), size=4.0, color=(1, 1, 1), spot=None, hidden=False):
    ld = D.lights.new(name, kind)
    ld.energy = energy
    ld.color = color
    if kind == 'AREA':
        ld.shape = 'DISK'
        ld.size = size
    elif kind == 'SUN':
        ld.angle = math.radians(8)
    elif kind == 'SPOT':
        ld.spot_size = math.radians(spot or 40)
        ld.spot_blend = 0.6
        ld.shadow_soft_size = 0.06
    o = D.objects.new(name, ld)
    coll.objects.link(o)
    o.rotation_euler = rot
    o.location = loc
    if hidden:
        o.visible_camera = False
        o.visible_glossy = False
        o.visible_transmission = False
    LIGHTS.append(o)
    return o


def clear_lights():
    for o in LIGHTS:
        ld = o.data
        D.objects.remove(o)
        D.lights.remove(ld)
    LIGHTS.clear()
    for o in [o for o in coll.objects if o.name.startswith("_haze")]:
        me = o.data
        D.objects.remove(o)
        D.meshes.remove(me)


def aim(loc, at, lens):
    cam.data.type = 'PERSP'
    cam.location = loc
    cam.rotation_euler = (Vector(at) - Vector(loc)).to_track_quat('-Z', 'Y').to_euler()
    cam.data.lens = lens
    cam.data.sensor_fit = 'AUTO'


def aim_ortho(centre, right, up, height_m):
    cam.data.type = 'ORTHO'
    fwd = up.cross(right)
    R = Matrix((right, up, -fwd)).transposed()
    cam.matrix_world = Matrix.Translation(centre - fwd * 40.0) @ R.to_4x4()
    cam.data.sensor_fit = 'VERTICAL'
    cam.data.ortho_scale = height_m


def shoot(path, w, h):
    sc.render.resolution_x, sc.render.resolution_y = w, h
    sc.render.filepath = path
    bpy.ops.render.render(write_still=True)
    print("SHOT", path, flush=True)


# ------------------------------------------------------------------ the two lighting setups
def studio():
    clear_lights()
    w = D.worlds.get("_studio") or D.worlds.new("_studio")
    w.use_nodes = True
    nt = w.node_tree
    for n in list(nt.nodes):
        if n.type != 'BACKGROUND' and n.type != 'OUTPUT_WORLD':
            nt.nodes.remove(n)
    bg = nt.nodes["Background"]
    tc = nt.nodes.new("ShaderNodeTexCoord")
    sx = nt.nodes.new("ShaderNodeSeparateXYZ")
    mr = nt.nodes.new("ShaderNodeMapRange")
    cr = nt.nodes.new("ShaderNodeValToRGB")
    nt.links.new(tc.outputs["Generated"], sx.inputs[0])
    nt.links.new(sx.outputs["Z"], mr.inputs["Value"])
    mr.inputs["From Min"].default_value = -1.0
    mr.inputs["From Max"].default_value = 1.0
    nt.links.new(mr.outputs["Result"], cr.inputs["Fac"])
    nt.links.new(cr.outputs["Color"], bg.inputs[0])
    cr.color_ramp.elements[0].color = (0.30, 0.32, 0.34, 1)
    cr.color_ramp.elements[1].color = (0.95, 0.97, 1.0, 1)
    bg.inputs[1].default_value = 1.0
    sc.world = w
    sc.view_settings.exposure = 0.0
    light("_key", 'SUN', 3.6, (math.radians(35), math.radians(-25), math.radians(-30)))
    light("_fill", 'AREA', 520, (math.radians(65), 0, math.radians(140)), (-6, 6, 3), size=8)
    light("_low", 'AREA', 200, (math.radians(150), 0, 0), (0, 0, -6), size=10)
    light("_side_s", 'AREA', 300, (0, math.radians(90), 0), (-9, 0.3, 0.6), size=8)
    light("_side_p", 'AREA', 200, (0, math.radians(-90), 0), (9, 0.3, 0.6), size=8)
    light("_cabin", 'AREA', 60, (math.radians(180), 0, 0), (0, -1.2, 0.95), size=0.8, color=(0.92, 0.97, 1.0), hidden=True)


def water():
    clear_lights()
    w = D.worlds.get("_water") or D.worlds.new("_water")
    w.use_nodes = True
    nt = w.node_tree
    for n in list(nt.nodes):
        if n.type != 'BACKGROUND' and n.type != 'OUTPUT_WORLD':
            nt.nodes.remove(n)
    bg = nt.nodes["Background"]
    tc = nt.nodes.new("ShaderNodeTexCoord")
    sx = nt.nodes.new("ShaderNodeSeparateXYZ")
    mr = nt.nodes.new("ShaderNodeMapRange")
    cr = nt.nodes.new("ShaderNodeValToRGB")
    nt.links.new(tc.outputs["Generated"], sx.inputs[0])
    nt.links.new(sx.outputs["Z"], mr.inputs["Value"])
    mr.inputs["From Min"].default_value = -1.0
    mr.inputs["From Max"].default_value = 1.0
    nt.links.new(mr.outputs["Result"], cr.inputs["Fac"])
    nt.links.new(cr.outputs["Color"], bg.inputs[0])
    cr.color_ramp.elements[0].color = (0.001, 0.004, 0.007, 1)
    cr.color_ramp.elements[1].color = (0.012, 0.05, 0.075, 1)
    bg.inputs[1].default_value = 1.0
    sc.world = w
    sc.view_settings.exposure = 0.35
    kl = light("_keyw", 'AREA', 5000, (0, 0, 0), (-7.5, -9.0, 6.5), size=5, color=(0.75, 0.90, 1.0))
    kl.rotation_euler = (Vector((0, -0.5, 0)) - kl.location).to_track_quat('-Z', 'Y').to_euler()
    light("_down", 'SUN', 0.9, (math.radians(15), math.radians(-10), 0), color=(0.55, 0.85, 1.0))
    light("_rim", 'AREA', 700, (math.radians(-65), 0, math.radians(200)), (-4, 8, 5), size=6, color=(0.45, 0.75, 1.0))
    light("_fillw", 'AREA', 260, (math.radians(70), 0, math.radians(-30)), (-6, -6, -2), size=6, color=(0.55, 0.85, 1.0))
    light("_cabin", 'AREA', 25, (math.radians(180), 0, 0), (0, -1.2, 0.95), size=0.8, color=(0.7, 0.92, 1.0), hidden=True)
    for n in ("SOCKET_Lamp_L", "SOCKET_Lamp_R", "SOCKET_PodLamp_L", "SOCKET_PodLamp_R", "SOCKET_DrumLamp_L", "SOCKET_DrumLamp_R"):
        if n not in D.objects:
            continue
        p = D.objects[n].matrix_world.translation + Vector((0, -0.05, 0))
        light("_spot_" + n, 'SPOT', 300 if "Drum" in n else 900, (math.radians(90), 0, math.radians(180)), tuple(p),
              color=(1.0, 0.95, 0.85), spot=46)
    p = D.objects["SOCKET_Scanner"].matrix_world.translation + Vector((0, -0.05, 0))
    light("_spot_scan", 'SPOT', 600, (math.radians(90), 0, math.radians(180)), tuple(p), color=(0.3, 0.7, 1.0), spot=70)
    bpy.ops.mesh.primitive_cube_add(size=60, location=(0, -6, 0))
    hz = bpy.context.active_object
    hz.name = "_haze"
    for c in hz.users_collection:
        c.objects.unlink(hz)
    coll.objects.link(hz)
    hm = D.materials.new("_haze")
    hm.use_nodes = True
    hn = hm.node_tree
    hn.nodes.remove(hn.nodes["Principled BSDF"])
    vs = hn.nodes.new("ShaderNodeVolumeScatter")
    vs.inputs["Color"].default_value = (0.35, 0.7, 0.85, 1)
    vs.inputs["Density"].default_value = 0.0035
    vs.inputs["Anisotropy"].default_value = 0.55
    hn.links.new(vs.outputs[0], hn.nodes["Material Output"].inputs["Volume"])
    hz.data.materials.append(hm)
    hz.visible_shadow = False


# ------------------------------------------------------------------ studio
if "studio" in MODES:
    studio()
    C = Vector((0.0, 0.25, -0.15))
    SHOTS = {"side": ((-15.0, 0.25, 0.2), C, 60), "top": ((-0.01, 0.25, 16.0), C, 50), "front": ((0.0, -15.0, 0.4), C, 95),
             "back": ((0.0, 15.0, 0.4), C, 95), "q_fore": ((-7.4, -8.6, 3.4), Vector((0, -0.4, -0.1)), 40),
             "q_aft": ((6.8, 8.2, 3.0), Vector((0, 0.6, -0.1)), 40)}
    for nm, (loc, at, lens) in SHOTS.items():
        if nm == "top":
            aim_ortho(Vector((0, 0.25, 0)), Vector((0, -1, 0)), Vector((1, 0, 0)), 5.4)
            cam.data.sensor_fit = 'HORIZONTAL'
            cam.data.ortho_scale = 8.6
            cam.location = Vector((0, 0.25, 20))
        else:
            aim(loc, at, lens)
        shoot(os.path.join(PREV, "%s_studio_%s.png" % (TAG, nm)), 1600, 1000)
    pose(hatch_open_deg=100.0)
    aim((0.35, 0.25, 2.55), Vector((0.0, -1.25, -0.35)), 18)
    shoot(os.path.join(PREV, "%s_studio_cabin_hatch.png" % TAG), 1400, 1000)
    pose(hatch_open_deg=0.0, door_r_deg=95.0)
    aim((-2.35, -0.05, 0.05), Vector((0.0, -1.0, -0.25)), 18)
    shoot(os.path.join(PREV, "%s_studio_cabin_door.png" % TAG), 1400, 1000)
    pose(door_r_deg=0.0)

# ------------------------------------------------------------------ water
if "water" in MODES:
    water()
    aim((4.6, -8.8, -2.2), Vector((0.0, -0.9, -0.3)), 32)
    shoot(os.path.join(PREV, "%s_water_nose.png" % TAG), 1600, 1000)
    aim((-6.4, 7.8, 2.6), Vector((0.0, 0.8, -0.1)), 34)
    shoot(os.path.join(PREV, "%s_water_stern.png" % TAG), 1600, 1000)

# ------------------------------------------------------------------ concept
if "concept" in MODES:
    water()
    c = json.load(open(os.path.join(HERE, "checks", "concept_cam.json")))
    cam.data.type = 'PERSP'
    cam.location = c["location"]
    cam.rotation_euler = Euler(c["rotation_euler_xyz_rad"], 'XYZ')
    cam.data.sensor_fit = 'HORIZONTAL'
    cam.data.sensor_width = 36.0
    cam.data.lens = c["lens_mm_36mm_sensor_horizontal"]
    shoot(os.path.join(CMP, "concept_render.png"), c["resolution"][0], c["resolution"][1])

# ------------------------------------------------------------------ lookalike (sheet framing, orthographic)
if "lookalike" in MODES:
    studio()
    sc.render.film_transparent = True
    sc.render.image_settings.color_mode = 'RGBA'
    K = 2
    VIEWS = {
        "side": ((430, 255, 1010, 520), 1.0, lambda cx, cy: Vector((0.0, G["Yi"](cx), G["Zi"](cy))), Vector((0, -1, 0)), Vector((0, 0, 1))),
        "top": ((440, 20, 1000, 250), 1.0, lambda cx, cy: Vector((G["Xt"](cy), G["Yi"](cx), 0.0)), Vector((0, -1, 0)), Vector((1, 0, 0))),
        "bottom": ((440, 520, 1000, 745), 1.0, lambda cx, cy: Vector(((cy - G["YC_BOTTOM"]) * S, G["Yi"](cx), 0.0)), Vector((0, -1, 0)), Vector((-1, 0, 0))),
        "front": ((95, 215, 385, 550), KF, lambda cx, cy: Vector((G["Xf"](cx), 0.0, G["Zi"](cy))), Vector((1, 0, 0)), Vector((0, 0, 1))),
        "back": ((1075, 235, 1345, 550), KB, lambda cx, cy: Vector((G["Xb"](cx), 0.0, G["Zi"](cy))), Vector((-1, 0, 0)), Vector((0, 0, 1))),
    }
    for nm, ((x0, y0, x1, y1), kh, centre, right, up) in VIEWS.items():
        w_px, h_px = x1 - x0, y1 - y0
        rw = int(round(w_px * K * kh))                 # the model's own uniform scale (gr_pairs.py stretches FRONT/BACK)
        rh = h_px * K
        aim_ortho(centre(0.5 * (x0 + x1), 0.5 * (y0 + y1)), right, up, h_px * S)
        shoot(os.path.join(CMP, "lk_%s.png" % nm), rw, rh)
    # SECTION A-A: the cut at the clamp axis (Y -0.09) looking forward; framed like the sheet's top-right section
    x0, y0, x1, y1 = 1040, 20, 1360, 240
    zc = G["Zs"](0.5 * (y0 + y1))
    xc = G["Xs"](0.5 * (x0 + x1))
    cam.data.type = 'ORTHO'
    cam.data.sensor_fit = 'VERTICAL'
    cam.data.ortho_scale = (y1 - y0) * G["SEC_KY"] * S
    cam.matrix_world = Matrix.Translation(Vector((xc, 10.0, zc))) @ Matrix(((-1, 0, 0), (0, 0, 1), (0, 1, 0))).transposed().to_4x4()
    cam.data.clip_start = 10.0 - (-0.09)
    cam.data.clip_end = 30.0
    # the cut removes the aft hull from the camera only (it still shades the cabin): light the cabin from inside
    light("_sec_top", 'AREA', 70, (0, 0, 0), (0.0, -0.45, 0.80), size=1.0, color=(0.95, 0.97, 1.0), hidden=True)
    light("_sec_pt", 'POINT', 18, (0, 0, 0), (0.0, -0.35, -0.2), color=(0.95, 0.97, 1.0), hidden=True)
    rw = int(round((x1 - x0) * K * G["SEC_KX"] / G["SEC_KY"]))
    shoot(os.path.join(CMP, "lk_section.png"), rw, (y1 - y0) * K)
    sc.render.film_transparent = False
    sc.render.image_settings.color_mode = 'RGB'
    cam.data.ortho_scale = 3.1
    cam.matrix_world = Matrix.Translation(Vector((0.0, 10.0, -0.05))) @ Matrix(((-1, 0, 0), (0, 0, 1), (0, 1, 0))).transposed().to_4x4()
    shoot(os.path.join(CMP, "lk_section_close.png"), 1100, 1100)
    cam.data.clip_start = 0.02
    sc.render.film_transparent = False
    sc.render.image_settings.color_mode = 'RGB'
print("GROUPER RENDERS OK", TAG, MODES)