"""Create an original native source clip and a frozen-event camera comparison with no model execution."""

import json
import math
import sys
from pathlib import Path

import bpy
import numpy as np
from bpy_extras.object_utils import world_to_camera_view
from mathutils import Vector

out = Path(sys.argv[sys.argv.index("--") + 1]).resolve()
out.mkdir()
for name in ["source-frames", "target-frames"]:
    (out / name).mkdir()
bpy.ops.object.select_all(action="SELECT")
bpy.ops.object.delete(use_global=False)
s = bpy.context.scene
s.render.engine = "BLENDER_EEVEE"
s.render.resolution_x = 832
s.render.resolution_y = 464
s.render.resolution_percentage = 100
s.render.fps = 16
s.render.image_settings.file_format = "PNG"
s.world.color = (0.15, 0.15, 0.15)
s.view_settings.view_transform = "Standard"
Kref = np.array([415.5298, 415.6922, 415.77786, 239.77779], dtype=np.float64)
K = Kref.copy()
K[[1, 3]] *= 464 / 480
ratio = K[0] / K[1]
s.render.pixel_aspect_x = 1
s.render.pixel_aspect_y = float(ratio)


def mat(name, color):
    m = bpy.data.materials.new(name)
    m.use_nodes = True
    node = m.node_tree.nodes.get("Principled BSDF")
    node.inputs["Base Color"].default_value = (*color, 1)
    node.inputs["Roughness"].default_value = 0.6
    return m


blue = mat("Blue arch", (0.08, 0.27, 0.48))
orange = mat("Orange ball", (0.95, 0.34, 0.08))
teal = mat("Teal floor", (0.18, 0.45, 0.42))
cream = mat("Cream marker", (0.8, 0.74, 0.53))


def cube(name, pos, dims, material):
    bpy.ops.mesh.primitive_cube_add(size=1, location=pos)
    obj = bpy.context.object
    obj.name = name
    obj.dimensions = dims
    bpy.ops.object.transform_apply(location=False, rotation=False, scale=True)
    obj.data.materials.append(material)
    return obj


cube("Floor", (0, 1, -0.12), (10, 10, 0.24), teal)
for x in [-0.9, 0.9]:
    cube(f"Arch post {x}", (x, 1.5, 1), (0.3, 0.4, 2), blue)
cube("Arch lintel", (0, 1.5, 2.05), (2.1, 0.4, 0.3), blue)
cube("Near marker", (-1.8, -0.4, 0.25), (0.4, 0.4, 0.5), cream)
bpy.ops.mesh.primitive_uv_sphere_add(segments=32, ring_count=16, radius=0.4)
ball = bpy.context.object
ball.name = "Original orange ball"
ball.data.materials.append(orange)
bpy.ops.object.shade_smooth()
bpy.ops.object.camera_add(location=(0, -6, 2.8))
camera = bpy.context.object
camera.name = "Calibrated source and target camera"
camera.rotation_euler = (
    (Vector((0, 0.7, 1.0)) - camera.location).to_track_quat("-Z", "Y").to_euler()
)
camera.data.sensor_fit = "HORIZONTAL"
camera.data.sensor_width = 36
camera.data.lens = float(K[0] * 36 / 832)
camera.data.shift_x = float((416 - K[2]) / 832)
camera.data.shift_y = float((K[3] - 232) * ratio / 832)
s.camera = camera
bpy.ops.object.light_add(type="AREA", location=(-3, -3, 6))
light = bpy.context.object
light.data.energy = 1100
light.data.size = 5
light.rotation_euler = (Vector((0, 0.6, 0.7)) - light.location).to_track_quat("-Z", "Y").to_euler()
conversion = np.diag([1, -1, -1, 1])
bpy.context.view_layer.update()
source_cv = np.array(camera.matrix_world) @ conversion


def ball_pos(index):
    return (-1.5 + 3 * index / 84, 0.6, 0.55 + 0.4 * math.sin(math.pi * index / 84))


def projection_check():
    cv = np.array(camera.matrix_world) @ conversion
    point = np.array([*ball.location, 1])
    local = np.linalg.inv(cv) @ point
    pixel = [K[0] * local[0] / local[2] + K[2], K[1] * local[1] / local[2] + K[3]]
    native = world_to_camera_view(s, camera, Vector(point[:3]))
    native_px = [native.x * 832, (1 - native.y) * 464]
    assert max(abs(a - b) for a, b in zip(pixel, native_px)) < 0.001, (pixel, native_px)
    return cv, pixel, native_px


s.frame_start, s.frame_end = 1, 85
for i in range(85):
    ball.location = ball_pos(i)
    ball.keyframe_insert(data_path="location", frame=i + 1)
    s.frame_set(i + 1)
    bpy.context.view_layer.update()
    projection_check()
    s.render.filepath = str(out / "source-frames" / f"frame-{i:03d}.png")
    bpy.ops.render.render(write_still=True)
s.frame_set(1)
bpy.ops.wm.save_as_mainfile(filepath=str(out / "wiw-source-study.blend"))
ball.animation_data_clear()
remap = list(range(40)) + [40] * 16 + list(range(40, 85))
assert len(remap) == 101 and remap.count(40) == 17
poses = []
records = []
s.frame_end = 93
for i in range(93):
    x = 0.75 * math.sin(math.pi * (i - 40) / 15) if 40 <= i <= 55 else 0
    camera.location.x = x
    camera.keyframe_insert(data_path="location", frame=i + 1)
    ball.location = ball_pos(remap[i])
    ball.keyframe_insert(data_path="location", frame=i + 1)
    s.frame_set(i + 1)
    bpy.context.view_layer.update()
    cv, pixel, native_px = projection_check()
    relative = np.linalg.inv(source_cv) @ cv
    poses.append(relative)
    records.append(
        {
            "frame": i,
            "sourceFrame": remap[i],
            "cameraX": x,
            "ballWorld": list(ball.location),
            "projectedPixels": pixel,
            "nativePixels": native_px,
        }
    )
    s.render.filepath = str(out / "target-frames" / f"frame-{i:03d}.png")
    bpy.ops.render.render(write_still=True)
poses = np.stack(poses).astype(np.float32)
assert np.allclose(poses[0], np.eye(4), atol=1e-6)
assert np.allclose(poses[-4:], poses[-1], atol=1e-6)
np.save(out / "poses.npy", poses, allow_pickle=False)
np.save(out / "intrinsics.npy", np.tile(Kref, (93, 1)).astype(np.float32), allow_pickle=False)
s.frame_set(1)
bpy.ops.wm.save_as_mainfile(filepath=str(out / "wiw-freeze-study.blend"))
record = {
    "blender": bpy.app.version_string,
    "width": 832,
    "height": 464,
    "fps": 16,
    "baseSourceFrames": 85,
    "freezeFrame": 40,
    "insertedFrames": 16,
    "preparedSourceFrames": 101,
    "plannedGeneratedFrames": 93,
    "nativeComparisonFrames": 93,
    "remap": remap,
    "intrinsicsReference832x480": Kref.tolist(),
    "workingIntrinsics832x464": K.tolist(),
    "sourceCameraToWorldOpenCV": source_cv.tolist(),
    "poseConvention": "OpenCV camera-to-world relative to the source camera at frame zero",
    "frames": records,
    "scope": "Original native source and re-render only; no depth estimation, upstream preparation, latent encoding or WorldinWorld inference",
}
(out / "wiw-native-measurements.json").write_text(json.dumps(record, indent=2) + "\n")
print("WIW_NATIVE_STUDY_COMPLETE", out)