"""Project a reference onto visible UV texels; retain occluded and shared-backside texels."""
import argparse
import json
from pathlib import Path
import sys

import bpy
import numpy as np
from mathutils import Vector
from mathutils.bvhtree import BVHTree

p = argparse.ArgumentParser()
p.add_argument('--folder', required=True)
p.add_argument('--front-axis', choices=['-Y', '+Y', '-X', '+X'], default='-Y')
p.add_argument('--strength', type=float, default=.9)
a = p.parse_args(sys.argv[sys.argv.index('--') + 1:])
folder = Path(a.folder)
bpy.ops.object.select_all(action='SELECT')
bpy.ops.object.delete(use_global=False)
bpy.ops.import_scene.gltf(filepath=str(folder / 'source.glb'))
objects = [o for o in bpy.context.scene.objects if o.type == 'MESH']
front = np.array({'-Y': [0, -1, 0], '+Y': [0, 1, 0], '-X': [-1, 0, 0], '+X': [1, 0, 0]}[a.front_axis], dtype=float)
right = np.cross(front, [0, 0, 1]) * -1
vertices, triangles, records = [], [], []
offset = 0
for obj in objects:
    mesh = obj.data
    mesh.calc_loop_triangles()
    world = np.array([obj.matrix_world @ v.co for v in mesh.vertices])
    normals = np.array([obj.matrix_world.to_3x3().inverted().transposed() @ v.normal for v in mesh.vertices])
    normals /= np.maximum(np.linalg.norm(normals, axis=1, keepdims=True), 1e-12)
    vertices.extend(world.tolist())
    triangles.extend(tuple(offset + i for i in t.vertices) for t in mesh.loop_triangles)
    records.append((obj, world, normals))
    offset += len(world)
vertices = np.array(vertices)
bvh = BVHTree.FromPolygons(vertices.tolist(), triangles, all_triangles=True)
height = np.ptp(vertices[:, 2])
horizontal = vertices @ right
span = np.ptp(horizontal)
horizontal_min = horizontal.min()
vertical_min = vertices[:, 2].min()
if height <= 0 or span <= 0:
    raise ValueError('Degenerate projection bounds')
reference = np.load(folder / 'reference.npz')
rgba, mask, bounds = reference['rgba'], reference['mask'], reference['bounds']
ref_h, ref_w = mask.shape
x0, y0, x1, y1 = bounds
metrics = {'front_axis': a.front_axis, 'strength': a.strength,
           'method': 'per-texel UV projection with BVH occlusion, normal falloff and shared-UV protection',
           'material_color_boundaries_protected': True, 'geometry_unchanged': True, 'uv_unchanged': True, 'pbr_channels_preserved': True, 'parts': []}
inputs = json.loads((folder / 'inputs.json').read_text())
for item in inputs:
    matches = [r for r in records if r[0].name == item['part_id'] or r[0].name.startswith(item['part_id'] + ' ')]
    if len(matches) != 1:
        raise ValueError(f"Cannot identify unique component {item['part_id']}")
    obj, world, normals = matches[0]
    source = np.load(folder / item['file'])
    h, w = source.shape[:2]
    weights = np.zeros((h, w), dtype=np.float32)
    colors = np.zeros((h, w, 3), dtype=np.uint8)
    blocked = np.zeros((h, w), dtype=bool)
    uv_data = obj.data.uv_layers.active.data
    for triangle in obj.data.loop_triangles:
        uv = np.array([uv_data[i].uv for i in triangle.loops])
        if np.any(uv < -1e-6) or np.any(uv > 1 + 1e-6):
            raise ValueError('Tiled UVs are unsupported for reference projection; source retained.')
        pixel = uv * [w, -h] + [0, h]
        low = np.maximum(np.floor(pixel.min(0)).astype(int), 0)
        high = np.minimum(np.ceil(pixel.max(0)).astype(int), [w - 1, h - 1])
        if np.any(high < low):
            continue
        yy, xx = np.mgrid[low[1]:high[1] + 1, low[0]:high[0] + 1]
        points = np.stack([xx.ravel() + .5, yy.ravel() + .5], axis=1)
        v0, v1 = pixel[1] - pixel[0], pixel[2] - pixel[0]
        determinant = v0[0] * v1[1] - v0[1] * v1[0]
        if abs(determinant) < 1e-10:
            continue
        delta = points - pixel[0]
        b1 = (delta[:, 0] * v1[1] - delta[:, 1] * v1[0]) / determinant
        b2 = (v0[0] * delta[:, 1] - v0[1] * delta[:, 0]) / determinant
        bary = np.stack([1 - b1 - b2, b1, b2], axis=1)
        inside = np.all(bary >= -1e-7, axis=1)
        if not inside.any():
            continue
        bary = bary[inside]
        xx, yy = xx.ravel()[inside], yy.ravel()[inside]
        ids = list(triangle.vertices)
        xyz = bary @ world[ids]
        n = bary @ normals[ids]
        facing = np.clip((n @ front - .35) / .55, 0, 1)
        facing = facing * facing * (3 - 2 * facing)
        rx = np.clip(np.rint(x0 + (xyz @ right - horizontal_min) / span * (x1 - x0)).astype(int), 0, ref_w - 1)
        ry = np.clip(np.rint(y1 - (xyz[:, 2] - vertical_min) / height * (y1 - y0)).astype(int), 0, ref_h - 1)
        eligible = (facing > 0) & mask[ry, rx]
        visible = np.zeros(len(xyz), dtype=bool)
        for index in np.flatnonzero(eligible):
            point = xyz[index]
            hit, _, _, _ = bvh.ray_cast(Vector(point + front * height * 2), Vector(-front))
            visible[index] = hit is not None and np.linalg.norm(np.array(hit) - point) < height * 1e-4
        original_rgb = source[yy, xx, :3].astype(float) / 255
        target_rgb = rgba[ry, rx, :3].astype(float) / 255
        original_sat = np.ptp(original_rgb, axis=1) / np.maximum(original_rgb.max(1), 1e-6)
        target_sat = np.ptp(target_rgb, axis=1) / np.maximum(target_rgb.max(1), 1e-6)
        colored = np.clip((original_sat - .35) / .15, 0, 1)
        washed_out = np.clip((.45 - target_sat) / .15, 0, 1)
        color_agreement = 1 - colored * washed_out
        source_chroma = original_rgb - original_rgb.mean(1, keepdims=True)
        target_chroma = target_rgb - target_rgb.mean(1, keepdims=True)
        cosine = np.sum(source_chroma * target_chroma, axis=1) / np.maximum(
            np.linalg.norm(source_chroma, axis=1) * np.linalg.norm(target_chroma, axis=1), 1e-6)
        changed_hue = colored * np.clip((target_sat - .35) / .15, 0, 1)
        color_agreement *= 1 - changed_hue * (1 - np.clip((cosine - .3) / .5, 0, 1))
        weight = facing * visible * color_agreement * a.strength
        blocked[yy[weight == 0], xx[weight == 0]] = True
        keep = weight > weights[yy, xx]
        weights[yy[keep], xx[keep]] = weight[keep]
        colors[yy[keep], xx[keep]] = rgba[ry[keep], rx[keep], :3]
    weights[blocked] = 0
    changed = weights > 0
    result = source.copy()
    result[:, :, :3] = np.rint(source[:, :, :3] * (1 - weights[:, :, None]) + colors * weights[:, :, None]).astype(np.uint8)
    np.save(folder / ('projected-' + item['file']), result)
    assert np.array_equal(source[~changed], result[~changed])
    assert np.array_equal(source[:, :, 3], result[:, :, 3])
    metrics['parts'].append({'part_id': item['part_id'], 'changed_texels': int(changed.sum()),
                             'total_texels': int(h * w), 'protected_texels_unchanged': True,
                             'alpha_unchanged': True})
(folder / 'projection.json').write_text(json.dumps(metrics, indent=2))
print(json.dumps(metrics), flush=True)