"""Render an original nested XYZ gimbal with measured native rotation matrices."""

import argparse
import json
import math
import sys
from pathlib import Path

import bpy
from mathutils import Euler

p = argparse.ArgumentParser()
p.add_argument("directory", type=Path)
a = p.parse_args(sys.argv[sys.argv.index("--") + 1 :])
r = a.directory
if r.exists():
    raise SystemExit("Choose a new output directory")
r.mkdir(parents=True)
bpy.ops.object.select_all(action="SELECT")
bpy.ops.object.delete(use_global=False)
s = bpy.context.scene
s.render.engine = "CYCLES"
s.cycles.samples = 16
s.cycles.use_denoising = True
s.render.resolution_x = 800
s.render.resolution_y = 800
s.render.resolution_percentage = 100
s.render.pixel_aspect_x = s.render.pixel_aspect_y = 1
s.render.image_settings.file_format = "PNG"
s.world.color = (0.05, 0.05, 0.05)


def material(name, c):
    m = bpy.data.materials.new(name)
    m.diffuse_color = (*c, 1)
    m.use_nodes = True
    n = m.node_tree.nodes.get("Principled BSDF")
    n.inputs["Base Color"].default_value = (*c, 1)
    n.inputs["Roughness"].default_value = 0.3
    return m


colors = [
    material("Inner X coral", (0.9, 0.18, 0.09)),
    material("Middle Y green", (0.2, 0.7, 0.3)),
    material("Outer Z blue", (0.12, 0.42, 0.9)),
]
cream = material("Original camera ivory", (0.8, 0.72, 0.53))
rig = []
for i, name in enumerate(["Outer Z", "Middle Y", "Inner X"]):
    o = bpy.data.objects.new(name, None)
    bpy.context.collection.objects.link(o)
    o.rotation_mode = "XYZ"
    if rig:
        o.parent = rig[-1]
    rig.append(o)
for i, (parent, radius, rot) in enumerate(
    [
        (rig[0], 1.8, (0, 0, 0)),
        (rig[1], 1.5, (math.pi / 2, 0, 0)),
        (rig[2], 1.2, (0, math.pi / 2, 0)),
    ]
):
    bpy.ops.mesh.primitive_torus_add(
        major_radius=radius, minor_radius=0.055, major_segments=96, minor_segments=12
    )
    o = bpy.context.object
    o.name = parent.name + " ring"
    o.parent = parent
    o.rotation_euler = rot
    o.data.materials.append(colors[2 - i])
bpy.ops.mesh.primitive_cube_add(size=1)
body = bpy.context.object
body.name = "Original camera body"
body.scale = (0.45, 0.3, 0.28)
body.parent = rig[2]
body.data.materials.append(cream)
b = body.modifiers.new("Soft edges", "BEVEL")
b.width = 0.08
b.segments = 3
bpy.ops.mesh.primitive_cylinder_add(vertices=64, radius=0.19, depth=0.3)
lens = bpy.context.object
lens.name = "Original lens"
lens.parent = rig[2]
lens.location = (0, -0.4, 0)
lens.rotation_euler = (math.pi / 2, 0, 0)
lens.data.materials.append(colors[2])
rig[0].rotation_euler.z = math.radians(30)
rig[2].rotation_euler.x = math.radians(20)
for frame, y in [(1, 0), (49, 90)]:
    rig[1].rotation_euler.y = math.radians(y)
    rig[1].keyframe_insert(data_path="rotation_euler", frame=frame)
s.frame_start = 1
s.frame_end = 49
s.render.fps = 24
bpy.ops.object.camera_add(location=(5, -7, 5))
cam = bpy.context.object
cam.rotation_euler = (-cam.location).to_track_quat("-Z", "Y").to_euler()
cam.data.type = "ORTHO"
cam.data.ortho_scale = 5.4
s.camera = cam
for loc, power, size in [((1, -4, 6), 850, 5), ((-4, 1, 3), 700, 4)]:
    bpy.ops.object.light_add(type="AREA", location=loc)
    o = bpy.context.object
    o.data.energy = power
    o.data.size = size
    o.rotation_euler = (-o.location).to_track_quat("-Z", "Y").to_euler()
records = []
for y in [0, 75, 90, 105]:
    rot = Euler(tuple(map(math.radians, [20, y, 30])), "XYZ")
    M = rot.to_matrix()
    q = rot.to_quaternion()
    records.append(
        {
            "degrees": {"x": 20, "y": y, "z": 30},
            "matrix": [list(row) for row in M],
            "quaternionWXYZ": list(q),
        }
    )
for frame, name in [(1, "open"), (49, "locked")]:
    s.frame_set(frame)
    bpy.context.view_layer.update()
    actual = rig[2].matrix_world.to_3x3()
    expected = Euler(tuple(map(math.radians, [20, 0 if frame == 1 else 90, 30])), "XYZ").to_matrix()
    assert max(abs(actual[i][j] - expected[i][j]) for i in range(3) for j in range(3)) < 1e-6
    s.render.filepath = str(r / f"gimbal-native-{name}.png")
    bpy.ops.render.render(write_still=True)
s.frame_set(1)
bpy.ops.wm.save_as_mainfile(filepath=str(r / "gimbal-native-study.blend"))
(r / "gimbal-native-measurements.json").write_text(
    json.dumps(
        {
            "Blender": bpy.app.version_string,
            "scope": "Original native rotation geometry not AI motion or physical stabilizer testing",
            "records": records,
        },
        indent=2,
    )
)