"""Preview renders (Cycles). Opens the .blend, adds a throw-away studio, renders, does not save.

Run headless:  blender -b --python render_previews.py [-- lineup full bust hands back] [--file name.blend] [--samples 96] [--tag v2]
"""
import bpy, sys, os
import numpy as np
from mathutils import Vector

HERE = os.path.dirname(os.path.abspath(__file__)); OUT = os.path.dirname(HERE)
PREV = os.path.join(OUT, 'previews'); os.makedirs(PREV, exist_ok=True)
ARGS = sys.argv[sys.argv.index('--') + 1:] if '--' in sys.argv else []


def opt(name, default):
    return ARGS[ARGS.index(name) + 1] if name in ARGS else default


WANT = {a for a in ARGS if a in ('lineup', 'full', 'bust', 'hands', 'back')} or {'lineup', 'full', 'bust', 'hands', 'back'}
TAG = opt('--tag', ''); SAMPLES = int(opt('--samples', '96'))
bpy.ops.wm.open_mainfile(filepath=os.path.join(OUT, opt('--file', 'Convict_Body_Augments.blend')))
sc = bpy.context.scene; vl = bpy.context.view_layer; sc.frame_set(1)


def lc(name, node=None):
    node = node or vl.layer_collection
    if node.collection.name == name: return node
    for ch in node.children:
        r = lc(name, ch)
        if r: return r


for n in [c.name for c in bpy.data.collections if c.name == 'Body_Core' or c.name.startswith('Stage')]:
    if lc(n): lc(n).exclude = True                         # only the line-up instances are drawn
EMP = [bpy.data.objects[n] for n in sorted(o.name for o in bpy.data.objects if o.name.startswith('Lineup_Stage'))]
arm = bpy.data.objects['Armature']; arm.hide_render = True

# ---- studio
st = bpy.data.collections.new('Studio'); sc.collection.children.link(st)
w = bpy.data.worlds.new('StudioWorld'); sc.world = w; w.use_nodes = True; nt = w.node_tree; nt.nodes.clear()
tc = nt.nodes.new('ShaderNodeTexCoord'); sep = nt.nodes.new('ShaderNodeSeparateXYZ'); mr = nt.nodes.new('ShaderNodeMapRange')
cr = nt.nodes.new('ShaderNodeValToRGB'); bg = nt.nodes.new('ShaderNodeBackground'); ow = nt.nodes.new('ShaderNodeOutputWorld')
mr.inputs['From Min'].default_value = -1; mr.inputs['From Max'].default_value = 1
e = cr.color_ramp.elements; e[0].position = .32; e[0].color = (.008, .018, .020, 1); e[1].position = .95; e[1].color = (.26, .44, .44, 1)
m = cr.color_ramp.elements.new(.52); m.color = (.06, .12, .125, 1)
nt.links.new(tc.outputs['Generated'], sep.inputs[0]); nt.links.new(sep.outputs['Z'], mr.inputs['Value']); nt.links.new(mr.outputs[0], cr.inputs[0])
nt.links.new(cr.outputs[0], bg.inputs[0]); nt.links.new(bg.outputs[0], ow.inputs[0]); bg.inputs[1].default_value = 1.0

gm = bpy.data.materials.new('StudioFloor'); gm.use_nodes = True
b = next(n for n in gm.node_tree.nodes if n.type == 'BSDF_PRINCIPLED'); b.inputs['Base Color'].default_value = (.02, .035, .038, 1); b.inputs['Roughness'].default_value = .32
CX = sum(e_.location.x for e_ in EMP) / len(EMP)
bpy.ops.mesh.primitive_plane_add(size=80, location=(CX, 0, 0)); fl = bpy.context.active_object; fl.data.materials.append(gm)
for c in fl.users_collection: c.objects.unlink(fl)
st.objects.link(fl)
LIGHTS = []


def light(name, loc, target, power, size, col=(1, 1, 1)):
    ld = bpy.data.lights.new(name, 'AREA'); ld.energy = power; ld.size = size; ld.color = col
    o = bpy.data.objects.new(name, ld); st.objects.link(o); o['rel'] = tuple(loc); o['tgt'] = tuple(target); LIGHTS.append(o); return o


# offsets from the subject: a soft key, a low fill, two hard rims from behind, a top strip
light('Key', (-2.6, -3.6, 3.2), (0, 0, 1.1), 420, 3.0, (1, .96, .92)); light('Fill', (3.2, -2.6, 1.2), (0, 0, 1.0), 110, 4.0, (.7, .9, 1))
light('RimL', (-2.8, 3.0, 2.6), (0, 0, 1.2), 1100, 1.6, (.6, .95, 1)); light('RimR', (2.8, 3.0, 2.6), (0, 0, 1.2), 1100, 1.6, (.7, 1, .95))
light('Top', (0, -.4, 4.2), (0, 0, 1.0), 180, 3.0, (.9, 1, 1))
cd = bpy.data.cameras.new('Cam'); cam = bpy.data.objects.new('Cam', cd); st.objects.link(cam); sc.camera = cam

sc.render.engine = 'CYCLES'; cy = sc.cycles; cy.samples = SAMPLES; cy.use_denoising = True; cy.use_adaptive_sampling = True; cy.adaptive_threshold = .02
cy.max_bounces = 6; cy.device = 'CPU'
if '--gpu' in ARGS:
    try:
        pr = bpy.context.preferences.addons['cycles'].preferences; pr.compute_device_type = 'HIP'; pr.get_devices()
        for d in pr.devices: d.use = True
        cy.device = 'GPU'
    except Exception as ex: print('gpu', ex)
for look in ('AgX - Medium High Contrast', 'Medium High Contrast'):
    try: sc.view_settings.look = look; break
    except Exception: pass
sc.view_settings.exposure = 0.0
sc.render.image_settings.file_format = 'PNG'; sc.render.film_transparent = False


def shot(name, loc, target, lens, res, only=None, centre=None, flip=False):
    for e_ in EMP: e_.hide_render = only is not None and e_ is not only
    c = Vector((centre if centre is not None else CX, 0, 0)); s = -1 if flip else 1; wide = 1.0 if only is not None else 2.2
    for L in LIGHTS:
        r = Vector(L['rel']); L.location = c + Vector((r.x * wide, r.y * s, r.z)); t = c + Vector(L['tgt'])
        L.rotation_euler = (t - L.location).to_track_quat('-Z', 'Y').to_euler(); L.data.energy = L.data.energy
    cam.location = loc; cam.rotation_euler = (Vector(target) - Vector(loc)).to_track_quat('-Z', 'Y').to_euler(); cd.lens = lens
    sc.render.resolution_x, sc.render.resolution_y = res; sc.render.resolution_percentage = 100
    p = os.path.join(PREV, name + TAG + '.png'); sc.render.filepath = p; bpy.ops.render.render(write_still=True); print('RENDERED', p)
    return p


def sheet(name, paths):
    arrs = []
    for p in paths:
        im = bpy.data.images.load(p); wd, ht = im.size
        a = np.empty(wd * ht * 4, dtype=np.float32); im.pixels.foreach_get(a); arrs.append(a.reshape(ht, wd, 4)); bpy.data.images.remove(im)
    a = np.concatenate(arrs, axis=1); ht, wd = a.shape[:2]
    im = bpy.data.images.new(name, wd, ht, alpha=False); im.pixels.foreach_set(a.ravel()); im.file_format = 'PNG'
    im.filepath_raw = os.path.join(PREV, name + TAG + '.png'); im.save(); print('SHEET', im.filepath_raw)


if 'lineup' in WANT:
    for L in LIGHTS: L.data.energy *= 2.4
    shot('lineup_threequarter', (CX - 5.4, -8.2, 1.35), (CX, 0, .95), 62, (2400, 1260))
    if 'back' in WANT: shot('lineup_back', (CX + 3.2, 9.2, 1.35), (CX, 0, .95), 62, (2400, 1260), flip=True)
    for L in LIGHTS: L.data.energy /= 2.4
if 'full' in WANT:
    ps = [shot(f'stage{i + 1}_full', (e_.location.x - 2.2, -4.5, 1.05), (e_.location.x, 0, .95), 72, (760, 1300), only=e_, centre=e_.location.x) for i, e_ in enumerate(EMP)]
    sheet('stages_full', ps)
if 'bust' in WANT:
    ps = [shot(f'stage{i + 1}_bust', (e_.location.x - 1.0, -1.9, 1.5), (e_.location.x, 0, 1.36), 62, (760, 900), only=e_, centre=e_.location.x) for i, e_ in enumerate(EMP)]
    sheet('stages_bust', ps)
if 'hands' in WANT:
    dg = bpy.context.evaluated_depsgraph_get(); hp = arm.evaluated_get(dg).pose.bones['hand_l'].head.copy(); ps = []
    for i, e_ in enumerate(EMP):
        h = hp + e_.location
        ps.append(shot(f'stage{i + 1}_hand', h + Vector((.36, -.60, .04)), h + Vector((0, -.01, -.085)), 70, (640, 760), only=e_, centre=e_.location.x))
    sheet('stages_hand', ps)