"""Bake photo-guided vector outlines into the existing marker; run inside Blender.

Edits authored vector shapes, not the reference photograph. Existing stone and UVs
are reused. Color and tangent normal are baked for the same glTF material.
"""
import json
from pathlib import Path
import bpy
import numpy as np

ROOT = Path('D:/project/sunrise')
scene = bpy.data.scenes['Summit slab and marker']
bpy.context.window.scene = scene
body = bpy.data.objects['Photo matched summit boulder']
mat = body.data.materials[0]
nodes, links = mat.node_tree.nodes, mat.node_tree.links
bs = next(n for n in nodes if n.type == 'BSDF_PRINCIPLED')
atlas = next(n for n in nodes if n.type == 'UVMAP' and n.uv_map == 'Painted monument atlas')
spec = json.loads((ROOT / 'scripts/constants/marker-inscription.json').read_text(encoding='utf-8'))
N = 2048
mask = np.zeros((N, N), dtype=np.float32)

def atlas_point(x, z, vertical_offset):
    ax, az = spec['layout']['anchor']
    scale = spec['layout']['scale']
    x = ax+(x-ax)*scale
    z = az+(z-az)*scale+vertical_offset
    return ((.02+.46*(x+1.3)/2.6)*N, (.02+.96*(z+.2)/2.3)*N)

def polygon(points, value):
    p = np.array(points)
    lo = np.maximum(0, np.floor(p.min(axis=0)).astype(int))
    hi = np.minimum(N-1, np.ceil(p.max(axis=0)).astype(int))
    X, Y = np.meshgrid(np.arange(lo[0],hi[0]+1)+.5, np.arange(lo[1],hi[1]+1)+.5)
    inside = np.zeros(X.shape, dtype=bool)
    for a, b in zip(p, np.roll(p, -1, axis=0)):
        if abs(b[1]-a[1]) < 1e-9:
            continue
        inside ^= ((a[1]>Y) != (b[1]>Y)) & (X < (b[0]-a[0])*(Y-a[1])/(b[1]-a[1])+a[0])
    patch = mask[lo[1]:hi[1]+1,lo[0]:hi[0]+1]
    patch[inside] = value

for glyph_index, glyph in enumerate(spec['glyphs']):
    x0,y0,x1,y1 = glyph['bounds']
    cx,top,w,h = glyph['placement']
    for shape in glyph['shapes']:
        for i, ring in enumerate(shape):
            points = []
            for x,y in ring:
                px = cx+((x-x0)/(x1-x0)-.5)*w
                pz = top-(y-y0)/(y1-y0)*h
                points.append(atlas_point(px,pz,spec['layout']['glyphVerticalOffsets'][glyph_index]))
            polygon(points, 1 if i == 0 else 0)

# Rasterize small serif lettering into the same atlas; no font or text mesh is
# exported. Separate runs match the photo's smaller Korean label and unit.
caption = spec['caption']
font = bpy.data.fonts.load(caption['font'], check_existing=True)
left = caption['left']
for run in caption['runs']:
    curve = bpy.data.curves.new('Temporary elevation lettering', type='FONT')
    curve.body = run['text']
    curve.font = font
    curve.resolution_u = 16
    obj = bpy.data.objects.new('Temporary elevation lettering', curve)
    scene.collection.objects.link(obj)
    bpy.context.view_layer.update()
    evaluated = obj.evaluated_get(bpy.context.evaluated_depsgraph_get())
    mesh = evaluated.to_mesh()
    try:
        mesh.calc_loop_triangles()
        coordinates = np.array([(v.co.x,v.co.y) for v in mesh.vertices])
        low, high = coordinates.min(axis=0), coordinates.max(axis=0)
        for tri in mesh.loop_triangles:
            points = []
            for vertex in tri.vertices:
                x,y = (coordinates[vertex]-low)/(high-low)
                px = left+x*run['width']
                pz = caption['baseline']+y*run['height']
                points.append(atlas_point(px,pz,spec['layout']['captionVerticalOffset']))
            polygon(points,1)
    finally:
        evaluated.to_mesh_clear()
        bpy.data.objects.remove(obj,do_unlink=True)
        bpy.data.curves.remove(curve)
    left += run['width']+run['gap']
if font.users == 0:
    bpy.data.fonts.remove(font)

# A sub-texel soft edge preserves the chisel silhouette without stair steps.
soft = (mask*4 + np.roll(mask,1,0)+np.roll(mask,-1,0)+np.roll(mask,1,1)+np.roll(mask,-1,1))/8
rng = np.random.default_rng(194)
wear = rng.uniform(.94,1,mask.shape).astype(np.float32)
wear[rng.random(mask.shape)<.007] = .55

def write_image(name, values):
    image = bpy.data.images.get(name) or bpy.data.images.new(name,width=N,height=N,alpha=False)
    if tuple(image.size) != (N,N):
        image.scale(N,N)
    rgba = np.ones((N,N,4),dtype=np.float32)
    rgba[:,:,:3] = values[:,:,None]
    image.pixels.foreach_set(rgba.ravel())
    image.update()
    image.pack()
    return image

write_image('Hanja pigment mask', soft*wear)
paint = next(n for n in nodes if n.type == 'MIX_RGB' and n.inputs[0].is_linked)
paint.inputs[2].default_value = (.008,.012,.014,1)
target = next(n for n in nodes if n.type == 'TEX_IMAGE' and n.image and n.image.name == 'Painted Hanja on granite')
links.new(paint.outputs[0], bs.inputs['Base Color'])
nodes.active = target
for obj in scene.objects:
    obj.select_set(False)
body.select_set(True)
bpy.context.view_layer.objects.active = body
scene.render.bake.use_pass_direct = False
scene.render.bake.use_pass_indirect = False
scene.render.bake.use_pass_color = True
scene.render.bake.margin = 8
bake_types = bpy.ops.object.bake.get_rna_type().properties['type'].enum_items
bpy.ops.object.bake(type=next(i.identifier for i in bake_types if i.identifier == 'DIFFUSE'))
target.image.save()
target.image.pack()
links.new(target.outputs['Color'],bs.inputs['Base Color'])

# Bake the shallow engraving plus the original granite normal into the atlas.
# Reuse nodes on rerun; always start from the original granite normal.
def node(kind, label):
    found = next((n for n in nodes if n.label == label),None)
    if found is None:
        found = nodes.new(kind)
        found.label = label
    return found

depth = node('ShaderNodeTexImage','Inscription depth source')
depth.image = write_image('Inscription recessed height',1-soft)
depth.image.colorspace_settings.name = 'Non-Color'
links.new(atlas.outputs['UV'],depth.inputs['Vector'])
original_normal = next(n for n in nodes if n.type == 'NORMAL_MAP' and n.label != 'Baked inscription normal')
# The old material used a default tangent basis; explicitly retain its grain UV.
original_normal.uv_map = 'Granite grain UV'
bump = node('ShaderNodeBump','Shallow inscription bevel')
bump.inputs['Strength'].default_value = .65
bump.inputs['Distance'].default_value = .0025
links.new(depth.outputs['Color'],bump.inputs['Height'])
links.new(original_normal.outputs['Normal'],bump.inputs['Normal'])
links.new(bump.outputs['Normal'],bs.inputs['Normal'])
normal_target = node('ShaderNodeTexImage','Baked inscription normal texture')
normal_image = bpy.data.images.get('Marker inscription normal') or bpy.data.images.new('Marker inscription normal',width=N,height=N,alpha=False)
normal_image.colorspace_settings.name = 'Non-Color'
normal_target.image = normal_image
links.new(atlas.outputs['UV'],normal_target.inputs['Vector'])
nodes.active = normal_target
# The edge material needs its own temporary bake target to keep its maps intact.
temporary = []
for side in list(body.data.materials)[1:]:
    previous = side.node_tree.nodes.active
    t = side.node_tree.nodes.new('ShaderNodeTexImage')
    t.image = normal_image
    side.node_tree.nodes.active = t
    temporary.append((side,t,previous))
try:
    bpy.ops.object.bake(type=next(i.identifier for i in bake_types if i.identifier == 'NORMAL'))
finally:
    for side,t,previous in temporary:
        side.node_tree.nodes.remove(t)
        side.node_tree.nodes.active = previous
normal_image.filepath_raw = str(ROOT / 'public/terrain/materials/marker-inscription-normal.png')
normal_image.file_format = next(i.identifier for i in normal_image.bl_rna.properties['file_format'].enum_items if i.identifier == 'PNG')
normal_image.save()
normal_image.pack()
normal = node('ShaderNodeNormalMap','Baked inscription normal')
normal.uv_map = 'Painted monument atlas'
links.new(normal_target.outputs['Color'],normal.inputs['Color'])
links.new(normal.outputs['Normal'],bs.inputs['Normal'])
body['inscriptionStyle'] = 'Photo-guided silhouette polygons with negative spaces, dense charcoal pigment and shallow baked engraving'
scene.render.filepath = str(ROOT / 'assets/summit/blender-plateau-study/marker-detailed-inscription.png')
bpy.ops.render.render(write_still=True)
print('Detailed inscription baked:',int(np.count_nonzero(mask)),'texels, 3 glyphs +',caption['text'],'; color + tangent normal.')