"""Cycles pictures of the robot cell in a studio like the user's picture: a dark hall with a polished concrete floor,
a big soft key from the front-left and above, high overhead strips for the floor's reflections; inside, the cell's own
LED strips and downlight (plus area and spot lamps of the studio at their places, so the inside is lit as in the
picture). Standard view transform. The studio is saved in the file (collection "Studio", not exported).

    blender.exe -b ../RobotCell.blend --python rc_render.py -- [tag] [shots=picture,q_front_left,...] [samples=N]

previews/RobotCell_<tag>_<shot>.png; "picture" is the picture's own camera and framing (RC_PictureCamera).
"""
import bpy, bmesh, math, os, sys, json
from mathutils import Vector

HERE = os.path.dirname(os.path.abspath(__file__))
FOLDER = os.path.dirname(HERE)
OUTD = os.path.join(FOLDER, "previews")
argv = sys.argv[sys.argv.index("--") + 1:] if "--" in sys.argv else []
TAG = next((a for a in argv if "=" not in a), "v01")
SAMPLES = int(next((a.split("=")[1] for a in argv if a.startswith("samples=")), "160"))
KEY = float(next((a.split("=")[1] for a in argv if a.startswith("key=")), "1100"))
C = Vector((0.0, 0.0, 1.2))
V = lambda x, y, z: Vector((x, y, z))

# name: (camera, target, lens, resolution, frame, sliders)   ("picture" uses RC_PictureCamera)
SHOTS = {
    "picture": (None, None, None, (1200, 960), 1, {}),
    "q_front_left": (V(-5.6, -6.4, 2.6), V(0.0, 0.0, 1.15), 40, (1600, 1200), 1, {}),
    "q_back_right": (V(5.8, 5.2, 3.0), V(0.0, 0.0, 1.1), 40, (1600, 1200), 1, {}),
    "open": (V(3.6, -5.2, 1.75), V(-0.1, 0.0, 1.05), 32, (1600, 1200), 180, {"door_left": 1.0, "side_door": 1.0}),
    "robot": (V(0.9, -2.9, 1.55), V(-0.75, -0.05, 1.20), 35, (1600, 1200), 1, {"door_left": 1.0}),
    "tool": (V(1.45, -1.85, 1.25), V(0.25, -0.15, 0.95), 50, (1600, 1200), 180, {"door_right": 1.0}),
    "pod": (V(-0.75, -2.65, 1.45), V(-1.27, -1.55, 1.22), 50, (1200, 1200), 1, {}),
    "torch": (V(0.98, -0.82, 1.02), V(0.36, -0.10, 0.82), 50, (1600, 1200), 150, {"door_right": 1.0, "turntable": 230.0}),
}
only = next((a.split("=")[1].split(",") for a in argv if a.startswith("shots=")), None)
if only:
    SHOTS = {k: v for k, v in SHOTS.items() if k in only}


def studio():
    sc = bpy.context.scene
    col = bpy.data.collections.get("Studio")
    if col is None:
        col = bpy.data.collections.new("Studio")
        sc.collection.children.link(col)
    for o in list(col.objects):
        if o.name != "RC_PictureCamera":
            bpy.data.objects.remove(o)
    # the hall's floor: polished concrete, a big plane
    me = bpy.data.meshes.new("Studio_Floor")
    bm = bmesh.new()
    s = 40.0
    vs = [bm.verts.new(p) for p in ((-s, -s, 0), (s, -s, 0), (s, s, 0), (-s, s, 0))]
    bm.faces.new(vs)
    bm.to_mesh(me)
    bm.free()
    fl = bpy.data.objects.new("Studio_Floor", me)
    col.objects.link(fl)
    m = bpy.data.materials.get("Studio_Concrete") or bpy.data.materials.new("Studio_Concrete")
    m.use_nodes = True
    nt = m.node_tree
    nt.nodes.clear()
    b = nt.nodes.new('ShaderNodeBsdfPrincipled')
    o = nt.nodes.new('ShaderNodeOutputMaterial')
    nt.links.new(b.outputs[0], o.inputs["Surface"])
    nz = nt.nodes.new('ShaderNodeTexNoise'); nz.inputs["Scale"].default_value = 0.6; nz.inputs["Detail"].default_value = 8
    ramp = nt.nodes.new('ShaderNodeMapRange')
    ramp.inputs["To Min"].default_value = 0.028; ramp.inputs["To Max"].default_value = 0.060
    nt.links.new(nz.outputs["Fac"], ramp.inputs["Value"])
    nt.links.new(ramp.outputs[0], b.inputs["Base Color"])
    nz2 = nt.nodes.new('ShaderNodeTexNoise'); nz2.inputs["Scale"].default_value = 3.0
    r2 = nt.nodes.new('ShaderNodeMapRange')
    r2.inputs["To Min"].default_value = 0.12; r2.inputs["To Max"].default_value = 0.34
    nt.links.new(nz2.outputs["Fac"], r2.inputs["Value"])
    nt.links.new(r2.outputs[0], b.inputs["Roughness"])
    me.materials.append(m)

    def area(name, loc, size, power, tgt=C, shape='DISK', size_y=None, color=(1, 1, 1)):
        ld = bpy.data.lights.new(name, 'AREA')
        ld.shape = shape
        ld.size = size
        if size_y:
            ld.size_y = size_y
        ld.energy = power
        ld.color = color
        ob = bpy.data.objects.new(name, ld)
        col.objects.link(ob)
        ob.location = loc
        ob.rotation_euler = (Vector(tgt) - Vector(loc)).normalized().to_track_quat('-Z', 'Y').to_euler()
        return ob
    area("Studio_Key", C + V(-3.5, -6.5, 5.0), 5.0, KEY)
    area("Studio_Fill", C + V(5.5, -4.0, 1.0), 4.0, KEY * 0.18)
    area("Studio_Rim", C + V(-2.0, 6.0, 3.0), 4.0, KEY * 0.20)
    for k, (x, y) in enumerate(((-6, -2), (0, -4), (6, -2), (-6, 6), (0, 6), (6, 6), (-3, -12), (4, -12))):
        area("Studio_Hall_%d" % k, V(x, y, 9.0), 0.6, 500, tgt=V(x, y, 0), shape='RECTANGLE', size_y=2.4,
             color=(1.0, 0.97, 0.92))
    # the cell's own lights, as lamps (the strips are emissive too; these make the inside read as lit)
    body = bpy.data.objects["SM_RobotCell_Body"]
    led = [i for i, mm in enumerate(body.data.materials) if mm and mm.name == "M_RC_LED"]
    strips = {}
    for p in body.data.polygons:
        if p.material_index in led and p.normal.z < -0.9:
            c = body.matrix_world @ p.center
            strips[(round(c.x, 2), round(c.y, 2))] = (c, p.area)
    for k, (c, a) in enumerate(strips.values()):
        if a > 0.01:
            ob = area("Studio_Strip_%d" % k, c - V(0, 0, 0.003), 0.06, 24, tgt=c - V(0, 0, 1), shape='RECTANGLE',
                      size_y=0.55)
        else:
            ld = bpy.data.lights.new("Studio_Downlight", 'SPOT')
            ld.energy = 30
            ld.spot_size = math.radians(56)
            ld.spot_blend = 0.6
            ld.shadow_soft_size = 0.03
            ob = bpy.data.objects.new("Studio_Downlight", ld)
            col.objects.link(ob)
            ob.location = c - V(0, 0, 0.004)
            ob.rotation_euler = (V(0, 0.35, -1)).normalized().to_track_quat('-Z', 'Y').to_euler()
    w = sc.world or bpy.data.worlds.new("World")
    sc.world = w
    w.use_nodes = True
    w.node_tree.nodes["Background"].inputs["Color"].default_value = (0.020, 0.024, 0.030, 1)
    w.node_tree.nodes["Background"].inputs["Strength"].default_value = 1.0
    cam = bpy.data.objects.get("RC_View") or bpy.data.objects.new("RC_View", bpy.data.cameras.new("RC_View"))
    if cam.name not in col.objects:
        col.objects.link(cam)
    return cam


def setup(sc):
    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 = d.type == 'HIP'
        sc.cycles.device = 'GPU'
    except Exception as ex:
        print("GPU not set:", ex)
    sc.cycles.samples = SAMPLES
    sc.cycles.use_denoising = True
    sc.cycles.max_bounces = 8
    sc.cycles.transmission_bounces = 8
    sc.view_settings.view_transform = 'Standard'
    sc.view_settings.look = 'None'
    sc.view_settings.exposure = 0.0
    sc.render.image_settings.file_format = 'PNG'
    sc.render.image_settings.color_mode = 'RGB'
    sc.render.film_transparent = False


def main():
    sc = bpy.context.scene
    view = studio()
    setup(sc)
    for o in bpy.data.objects:
        if o.name.startswith("UCX_"):
            o.hide_render = True
    root = bpy.data.objects["RobotCell_Root"]
    act = root.animation_data.action if root.animation_data else None
    pic = bpy.data.objects["RC_PictureCamera"]
    done = []
    os.makedirs(OUTD, exist_ok=True)
    for nm, (loc, tgt, lens, res, frame, sliders) in SHOTS.items():
        root.animation_data.action = None if sliders else act
        for k in ("door_left", "door_right", "side_door", "turntable", "estop"):
            root[k] = sliders.get(k, 0.0)
        root.update_tag()
        sc.frame_set(frame + 1)
        sc.frame_set(frame)
        bpy.context.view_layer.update()
        sc.render.resolution_x, sc.render.resolution_y = res
        sc.render.resolution_percentage = 100
        if loc is None:
            sc.camera = pic
        else:
            sc.camera = view
            view.location = loc
            view.rotation_euler = (tgt - loc).normalized().to_track_quat('-Z', 'Y').to_euler()
            view.data.lens = lens
            view.data.sensor_fit = 'AUTO'
            view.data.clip_start = 0.02
        sc.render.filepath = os.path.join(OUTD, "RobotCell_%s_%s.png" % (TAG, nm))
        bpy.ops.render.render(write_still=True)
        done.append(sc.render.filepath)
    root.animation_data.action = act
    for k in ("door_left", "door_right", "side_door", "turntable", "estop"):
        root[k] = 0.0
    sc.camera = pic
    sc.frame_set(1)
    bpy.ops.wm.save_mainfile()
    print("RENDERS OK", done)


main()