# -*- coding: utf-8 -*-
"""Renders previews again from the built OrePatches.blend after textures were replaced on disk (the .blend links
them, so nothing else has to be rebuilt). Patches of the named ore types, their pieces, the mining stages and the sheet.

  blender -b ../OrePatches.blend --python rerender.py -- Ferrite [Iron ...]
"""
import bpy, os, sys, math
import numpy as np
from mathutils import Vector

HERE = os.path.dirname(os.path.abspath(__file__))
PRE = os.path.join(os.path.dirname(HERE), 'previews')
ORES = sys.argv[sys.argv.index('--') + 1:] if '--' in sys.argv else []
sc = bpy.context.scene
for im in bpy.data.images:
    if im.source == 'FILE' and not im.packed_file:
        im.reload()
prefs = bpy.context.preferences.addons['cycles'].preferences
for typ in ('HIP', 'OPTIX', 'CUDA'):
    try:
        prefs.compute_device_type = typ
        prefs.get_devices()
        if [d for d in prefs.devices if d.type == typ]:
            for d in prefs.devices:
                d.use = (d.type == typ)
            sc.cycles.device = 'GPU'
            break
    except Exception:
        pass
sc.cycles.samples, sc.cycles.use_denoising = 160, True
cam = bpy.data.objects['Studio_Camera']
lights = [bpy.data.objects['Studio_' + n] for n in ('Key', 'Rim', 'Fill')]
roots = sorted([o for o in bpy.data.objects if o.type == 'EMPTY' and o.name.startswith('OrePatch_')],
               key=lambda o: (round(o.location.x, 2), -round(o.location.y, 2)))        # the build's order: by ore, A before B
parts = {r.name: [c for c in r.children if c.type == 'MESH' and not c.name.startswith('UCX_')] for r in roots}
pieces = [o for o in bpy.data.objects if o.name.startswith(('SM_OreCluster_', 'SM_OreChunk_'))]
everything = [o for v in parts.values() for o in v] + pieces


def unit(v):
    v = np.asarray(v, float)
    return v / np.linalg.norm(v)


def aim(ob, target):
    ob.rotation_euler = (Vector(target) - ob.location).to_track_quat('-Z', 'Y').to_euler()


def frame(objs, side=1.0, view=None):                       # the same framing as build_ore_patches.py
    pts = []
    for ob in objs:
        mw = np.array(ob.matrix_world)
        V = np.empty(len(ob.data.vertices) * 3)
        ob.data.vertices.foreach_get('co', V)
        V = V.reshape(-1, 3) @ mw[:3, :3].T + mw[:3, 3]
        pts.append(V[V[:, 2] > -0.02])
    P = np.concatenate(pts)
    c = (P.max(0) + P.min(0)) * 0.5
    R = float(np.linalg.norm(P - c, axis=1).max())
    fov = 2 * math.atan(min(cam.data.sensor_width, cam.data.sensor_width * sc.render.resolution_y / sc.render.resolution_x) * 0.5 / cam.data.lens)
    d = R / math.sin(fov * 0.5) * 0.94
    if view is not None:
        d = (P[:, 0].max() - P[:, 0].min()) * 0.5 / math.tan(math.atan(cam.data.sensor_width * 0.5 / cam.data.lens)) * 1.12
    cam.location = Vector(c.tolist()) + Vector(unit(view or [-0.62 * side, -1.0, 0.50]).tolist()) * d
    aim(cam, c.tolist())
    for lo, off, pw in zip(lights, ((-1.4, -1.6, 2.4), (1.8, 1.9, 1.3), (2.2, -1.6, 0.7)), (260, 220, 50)):
        lo.location = Vector(c.tolist()) + Vector(off) * R * 1.5
        lo.data.energy = pw * R * R
        aim(lo, c.tolist())


def show(objs):
    for o in everything:
        o.hide_render = o not in objs


def render(name):
    sc.render.filepath = os.path.join(PRE, name)
    bpy.ops.render.render(write_still=True)
    print('[rerender]', name, flush=True)


def sheet(paths, cols, out):
    tiles = []
    for p in paths:
        im = bpy.data.images.load(p)
        w, h = im.size
        a = np.empty(w * h * 4, np.float32)
        im.pixels.foreach_get(a)
        bpy.data.images.remove(im)
        tiles.append(a.reshape(h, w, 4))
    h, w = tiles[0].shape[:2]
    rows = (len(tiles) + cols - 1) // cols
    big = np.zeros((rows * h, cols * w, 4), np.float32)
    big[..., 3] = 1
    for i, tl in enumerate(tiles):
        r, c = i // cols, i % cols
        big[(rows - 1 - r) * h:(rows - r) * h, c * w:(c + 1) * w] = tl
    im = bpy.data.images.new('_sheet', cols * w, rows * h, alpha=False)
    im.pixels.foreach_set(big.ravel())
    im.filepath_raw, im.file_format = out, 'PNG'
    im.save()
    bpy.data.images.remove(im)


sc.render.resolution_x, sc.render.resolution_y = 1280, 960
for i, r in enumerate(roots):
    if any(r.name.startswith('OrePatch_%s_' % o) for o in ORES):
        show(parts[r.name])
        frame(parts[r.name], side=1.0 if i % 2 == 0 else -1.0)
        render(r.name + '.png')
sheet([os.path.join(PRE, r.name + '.png') for r in roots], 4, os.path.join(PRE, 'OrePatches_sheet.png'))
r = roots[1]                                                # the mining stages are shown on the second patch
if any(r.name.startswith('OrePatch_%s_' % o) for o in ORES):
    insts = sorted([o for o in parts[r.name] if '_Ore_' in o.name], key=lambda o: o.name)
    rock = [o for o in parts[r.name] if '_Ore_' not in o.name]
    frame(parts[r.name])
    for k, keep in enumerate((len(insts), len(insts) // 2, 0)):
        show(rock + insts[:keep])
        render('%s_stage%d.png' % (r.name, k + 1))
    sheet([os.path.join(PRE, '%s_stage%d.png' % (r.name, k + 1)) for k in range(3)], 3, os.path.join(PRE, 'OrePatch_mining_stages.png'))
sc.render.resolution_x, sc.render.resolution_y = 2400, 500
for o in ORES:
    objs = [p for p in pieces if p.name.split('_')[2] == o]
    show(objs)
    frame(objs, view=[-0.12, -1.0, 0.42])
    render('OrePieces_%s.png' % o)
print('[rerender] done', ORES)