"""Cycles pictures of the nutrition vat in a studio like the user's picture (black room, a glossy dark floor, soft light
from the upper left, a rim light behind), Standard view transform so colours read as in the picture. Saves the studio
into the file (collection "Studio", not exported).

    blender.exe -b ../NutritionVat.blend --python nv_render.py -- [tag] [shots=fit_left,hero,...] [samples=N] [key=W]

fit_* shots use the cameras fitted to the picture's three views (build/cams.json) and its framing (twice its size,
only the view's panel rendered): previews/cycles/<tag>_<view>.png, laid beside the picture by
"python nv_compare.py <tag> --src cycles". The hoses show in fit_mid only (the picture's middle view has them).
Other shots: previews/NutritionVat_<tag>_<shot>.png.
"""
import bpy, math, os, sys, json
from mathutils import Vector, Euler

HERE = os.path.dirname(os.path.abspath(__file__))
FOLDER = os.path.dirname(HERE)
exec(open(os.path.join(HERE, "nv_lib.py")).read())
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=")), "96"))
KEY = float(next((a.split("=")[1] for a in argv if a.startswith("key=")), "170"))
cams = json.load(open(os.path.join(HERE, "cams.json")))
PANELS = {"left": (20, 20, 330, 570), "mid": (350, 10, 690, 570), "right": (670, 20, 990, 570)}

# name: (camera, target, lens, resolution, hoses)
SHOTS = {
    "fit_left": None, "fit_mid": None, "fit_right": None,
    "hero": (P(560, -60, 215), P(0, 500, 270), 50, (1600, 2000), False),
    "front_left": (P(-520, -80, 215), P(0, 500, 270), 50, (1600, 2000), False),
    "side_hoses": (P(980, 540, 270), P(0, 540, 280), 50, (1600, 2000), True),
    "rear": (P(560, 1120, 190), P(0, 520, 280), 50, (1600, 2000), False),
    "tank": (P(-110, 230, 230), P(0, 470, 250), 75, (1600, 2000), False),
    "top": (P(250, 300, -260), P(0, 495, 70), 60, (1600, 1200), False),
}
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):
        bpy.data.objects.remove(o)
    me = bpy.data.meshes.new("Studio_Floor")
    r = 30.0
    me.from_pydata([(-r, -r, 0), (r, -r, 0), (r, r, 0), (-r, r, 0)], [], [(0, 1, 2, 3)])
    floor = bpy.data.objects.new("Studio_Floor", me)
    col.objects.link(floor)
    fm = bpy.data.materials.get("Studio_Floor") or bpy.data.materials.new("Studio_Floor")
    fm.use_nodes = True
    b = fm.node_tree.nodes.get("Principled BSDF")
    b.inputs["Base Color"].default_value = (0.012, 0.013, 0.014, 1)
    b.inputs["Roughness"].default_value = 0.12
    me.materials.append(fm)
    centre = Vector((0.0, 0.0, 1.0))

    def area(name, loc, size, power, color=(1, 1, 1)):
        ld = bpy.data.lights.new(name, 'AREA')
        ld.shape = 'DISK'
        ld.size = size
        ld.energy = power
        ld.color = color
        o = bpy.data.objects.new(name, ld)
        col.objects.link(o)
        o.location = loc
        o.rotation_euler = (centre - loc).normalized().to_track_quat('-Z', 'Y').to_euler()
        return o
    area("Studio_Key", centre + Vector((-2.2, -2.6, 2.4)), 2.2, KEY)
    area("Studio_Fill", centre + Vector((3.0, -1.5, 0.6)), 2.5, KEY * 0.35)
    area("Studio_Rim", centre + Vector((1.5, 3.2, 1.8)), 1.6, KEY * 0.55, (0.85, 0.95, 1.0))
    area("Studio_Top", centre + Vector((0.0, 0.0, 3.2)), 1.5, KEY * 0.25)
    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.004, 0.004, 0.005, 1)
    w.node_tree.nodes["Background"].inputs["Strength"].default_value = 1.0
    cam = bpy.data.objects.new("Studio_Camera", bpy.data.cameras.new("Studio_Camera"))
    col.objects.link(cam)
    sc.camera = 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 = 10
    sc.cycles.transmission_bounces = 10
    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
    cam = studio()
    setup(sc)
    for o in bpy.data.objects:
        if o.name.startswith("UCX_"):
            o.hide_render = True
    hoses = bpy.data.objects.get("SM_NutritionVat_Hoses")
    done = []
    os.makedirs(os.path.join(OUTD, "cycles"), exist_ok=True)
    for nm, spec in SHOTS.items():
        cd = cam.data
        cd.shift_x = cd.shift_y = 0.0
        sc.render.use_border = False
        if nm.startswith("fit_"):
            v = nm[4:]
            c = cams[v]
            p = c["params"]
            cam.location = p[:3]
            cam.rotation_euler = Euler(p[3:6], 'XYZ')
            cd.sensor_fit = 'HORIZONTAL'
            cd.sensor_width = 36.0
            cd.lens = c["lens_mm"]
            cd.shift_x, cd.shift_y = c["shift"]
            dist = (cam.location - Vector((0, 0, 1.0))).length
            cd.clip_start, cd.clip_end = max(0.05, dist - 4.0), dist + 30.0
            sc.render.resolution_x, sc.render.resolution_y = 2000, 1158
            u0, v0, u1, v1 = PANELS[v]
            sc.render.use_border = True
            sc.render.use_crop_to_border = False
            sc.render.border_min_x, sc.render.border_max_x = u0 / 1000.0, u1 / 1000.0
            sc.render.border_min_y, sc.render.border_max_y = 1 - v1 / 579.0, 1 - v0 / 579.0
            show_h = v == "mid"
            path = os.path.join(OUTD, "cycles", "%s_%s.png" % (TAG, v))
        else:
            loc, tgt, lens, res, show_h = spec
            cam.location = loc
            cam.rotation_euler = (tgt - loc).normalized().to_track_quat('-Z', 'Y').to_euler()
            cd.sensor_fit = 'AUTO'
            cd.lens = lens
            cd.clip_start, cd.clip_end = 0.02, 100.0
            sc.render.resolution_x, sc.render.resolution_y = res
            path = os.path.join(OUTD, "NutritionVat_%s_%s.png" % (TAG, nm))
        if hoses:
            hoses.hide_render = not show_h
        sc.render.resolution_percentage = 100
        sc.render.filepath = path
        bpy.ops.render.render(write_still=True)
        done.append(os.path.basename(path))
    sc.render.use_border = False
    if hoses:
        hoses.hide_render = False
    bpy.ops.wm.save_mainfile()
    print("RENDERS OK", done)


main()