"""Preview renders of the four stages. Opens the .blend, adds a throw-away studio, renders, does not save.

Run headless:  blender -b --python render_previews.py [-- lineup hero hands back]
"""
import bpy, sys, os
import numpy as np
from math import radians
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 []
WANT = set(ARGS) or {'lineup', 'hero', 'hands', 'back'}
bpy.ops.wm.open_mainfile(filepath=os.path.join(OUT, '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 ('Body_Core', 'Stage1_Issue'): 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 = .30; e[0].color = (.010, .018, .020, 1); e[1].position = .95; e[1].color = (.30, .44, .44, 1)
m = cr.color_ramp.elements.new(.52); m.color = (.07, .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 = (.035, .055, .058, 1); b.inputs['Roughness'].default_value = .45
bpy.ops.mesh.primitive_plane_add(size=60, location=(2.9, 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)


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.location = loc
    o.rotation_euler = (Vector(target) - Vector(loc)).to_track_quat('-Z', 'Y').to_euler(); return o


CX = sum(e_.location.x for e_ in EMP) / len(EMP)
LIGHTS = [light('Key', (CX - 3.5, -5.5, 4.2), (CX, 0, 1.0), 1300, 5.0, (1, .97, .93)), light('Fill', (CX + 5, -4, 2.0), (CX, 0, 1.0), 450, 6.0, (.8, .95, 1)),
          light('RimL', (CX - 5, 4.5, 3.2), (CX, 0, 1.2), 3200, 3.0, (.7, .95, 1)), light('RimR', (CX + 5, 4.5, 3.2), (CX, 0, 1.2), 3200, 3.0, (.75, 1, .95)),
          light('Top', (CX, -1.0, 5.5), (CX, 0, 1.0), 600, 7.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 = 'BLENDER_EEVEE'
try:
    sc.eevee.taa_render_samples = 48; sc.eevee.use_raytracing = True
except Exception as ex: print('eevee opt', 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.35
sc.render.image_settings.file_format = 'PNG'; sc.render.film_transparent = False


def shot(name, loc, target, lens, res, only=None, flip_lights=False):
    for e_ in EMP: e_.hide_render = only is not None and e_ is not only
    cam.location = loc; cam.rotation_euler = (Vector(target) - Vector(loc)).to_track_quat('-Z', 'Y').to_euler(); cd.lens = lens
    for L in LIGHTS:
        if flip_lights and not L.get('f'): L.location.y *= -1; L['f'] = 1; L.rotation_euler = (Vector((CX, 0, 1.1)) - L.location).to_track_quat('-Z', 'Y').to_euler()
        if not flip_lights and L.get('f'): L.location.y *= -1; L['f'] = 0; L.rotation_euler = (Vector((CX, 0, 1.1)) - L.location).to_track_quat('-Z', 'Y').to_euler()
    sc.render.resolution_x, sc.render.resolution_y = res; sc.render.resolution_percentage = 100
    p = os.path.join(PREV, name + '.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 + '.png'); im.save(); print('SHEET', im.filepath_raw)


if 'lineup' in WANT:
    shot('lineup_front', (CX, -9.5, 1.05), (CX, 0, .93), 60, (2400, 1260))
    shot('lineup_threequarter', (CX - 5.2, -8.0, 1.5), (CX, 0, .93), 60, (2400, 1260))
if 'back' in WANT:
    shot('lineup_back', (CX + 3.0, 9.0, 1.4), (CX, 0, .93), 60, (2400, 1260), flip_lights=True)
if 'hero' in WANT:
    ps = []
    for i, e_ in enumerate(EMP):
        x = e_.location.x
        ps.append(shot(f'stage{i + 1}_full', (x - 2.1, -4.3, 1.2), (x, 0, .93), 70, (760, 1300), only=e_))
    sheet('stages_full', ps)
    ps = []
    for i, e_ in enumerate(EMP):
        x = e_.location.x
        ps.append(shot(f'stage{i + 1}_bust', (x - .95, -1.75, 1.5), (x, 0, 1.36), 60, (760, 900), only=e_))
    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((.30, -.52, .10)), h + Vector((0, 0, -.06)), 80, (640, 760), only=e_))
    sheet('stages_hand', ps)