"""Render an original native camera study and export OpenCV camera-to-world input arrays."""

import json
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()
(out / "frames").mkdir()
bpy.ops.object.select_all(action="SELECT")
bpy.ops.object.delete(use_global=False)
scene = bpy.context.scene
scene.render.engine = "BLENDER_EEVEE"
scene.render.resolution_x = 832
scene.render.resolution_y = 480
scene.render.resolution_percentage = 100
scene.render.pixel_aspect_x = 1
scene.render.pixel_aspect_y = 1
scene.render.fps = 16
scene.render.image_settings.file_format = "PNG"
scene.frame_start = 1
scene.frame_end = 81
scene.world.color = (0.16, 0.19, 0.22)
scene.view_settings.view_transform = "Standard"


def material(name, color):
    m = bpy.data.materials.new(name)
    m.diffuse_color = (*color, 1)
    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.7
    return m


teal = material("Teal floor", (0.17, 0.39, 0.4))
cream = material("Cream tower", (0.85, 0.77, 0.58))
coral = material("Coral roof", (0.68, 0.19, 0.13))
gold = material("Gold lamp", (0.96, 0.64, 0.12))
blue = material("Blue markers", (0.09, 0.27, 0.52))


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


cube("Floor", (0, 1, -0.12), (9, 10, 0.24), teal)
bpy.ops.mesh.primitive_cylinder_add(vertices=48, radius=0.45, depth=1.7, location=(0, 0.7, 0.85))
tower = bpy.context.object
tower.name = "Original beacon tower"
tower.data.materials.append(cream)
cube("Lamp", (0, 0.7, 1.91), (0.57, 0.57, 0.42), gold)
bpy.ops.mesh.primitive_cone_add(
    vertices=48, radius1=0.65, radius2=0, depth=0.55, location=(0, 0.7, 2.4)
)
bpy.context.object.name = "Beacon roof"
bpy.context.object.data.materials.append(coral)
for name, loc in [("Near marker", (-1.4, -0.5, 0.3)), ("Far marker", (1.4, 2.0, 0.3))]:
    cube(name, loc, (0.45, 0.45, 0.6), blue)
bpy.ops.object.camera_add(location=(0, -6, 2.7))
camera = bpy.context.object
camera.name = "Native translation camera"
camera.rotation_euler = (
    (Vector((0, 0.7, 1.1)) - camera.location).to_track_quat("-Z", "Y").to_euler()
)
camera.data.type = "PERSP"
camera.data.lens = 36
camera.data.sensor_width = 36
camera.data.sensor_fit = "HORIZONTAL"
camera.data.shift_x = 0
camera.data.shift_y = 0
scene.camera = camera
for frame in range(1, 82):
    camera.location.x = -0.8 + 1.6 * (frame - 1) / 80
    camera.keyframe_insert(data_path="location", frame=frame)
bpy.ops.object.light_add(type="AREA", location=(-3, -4, 7))
light = bpy.context.object
light.data.energy = 1200
light.data.size = 5
light.rotation_euler = (Vector((0, 0.7, 0.8)) - light.location).to_track_quat("-Z", "Y").to_euler()
poses = []
records = []
anchor = np.array([0.0, 0.7, 1.0, 1.0])
K = np.array([[832.0, 832.0, 416.0, 240.0]], dtype=np.float64)
conversion = np.diag([1.0, -1.0, -1.0, 1.0])
for frame in range(1, 82):
    scene.frame_set(frame)
    bpy.context.view_layer.update()
    pose = np.array(camera.matrix_world, dtype=np.float64) @ conversion
    poses.append(pose)
    pc = np.linalg.inv(pose) @ anchor
    assert pc[2] > 0
    projected = [832 * pc[0] / pc[2] + 416, 832 * pc[1] / pc[2] + 240]
    native = world_to_camera_view(scene, camera, Vector(anchor[:3]))
    native_px = [native.x * 832, (1 - native.y) * 480]
    assert max(abs(a - b) for a, b in zip(projected, native_px)) < 0.001
    records.append(
        {
            "frame": frame - 1,
            "seconds": (frame - 1) / 16,
            "cameraWorld": list(camera.location),
            "anchorPixels": projected,
            "nativeAnchorPixels": native_px,
        }
    )
    scene.render.filepath = str(out / "frames" / f"frame-{frame - 1:03d}.png")
    bpy.ops.render.render(write_still=True)
poses = np.stack(poses)
np.save(out / "poses.npy", poses, allow_pickle=False)
np.save(out / "intrinsics.npy", K, allow_pickle=False)
scene.frame_set(1)
bpy.ops.wm.save_as_mainfile(filepath=str(out / "lingbot-camera-study.blend"))
measurements = {
    "blender": bpy.app.version_string,
    "engine": scene.render.engine,
    "width": 832,
    "height": 480,
    "fps": 16,
    "frameCount": 81,
    "durationSeconds": 81 / 16,
    "lastTimestamp": 80 / 16,
    "cameraConvention": "OpenCV camera-to-world in unchanged Blender world coordinates",
    "cameraAxes": "+X right +Y down +Z forward",
    "conversion": "Blender camera matrix_world multiplied by diag(1,-1,-1,1)",
    "posesShape": list(poses.shape),
    "intrinsicsShape": list(K.shape),
    "intrinsicsOrder": ["fx", "fy", "cx", "cy"],
    "anchorWorld": anchor.tolist(),
    "frames": records,
    "scope": "Original native Blender inputs only with no LingBot inference",
}
(out / "lingbot-camera-measurements.json").write_text(json.dumps(measurements, indent=2) + "\n")
print("LINGBOT_CAMERA_STUDY_COMPLETE", out)