import csv
import json
import math
import sys
from pathlib import Path

import bpy
from bpy_extras.object_utils import world_to_camera_view
from mathutils import Vector

out = Path(sys.argv[sys.argv.index("--") + 1])
out.mkdir(parents=True, exist_ok=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 = 32
s.cycles.use_denoising = True
s.render.resolution_x = 640
s.render.resolution_y = 400
s.render.resolution_percentage = 100
s.render.image_settings.file_format = "PNG"
s.render.image_settings.color_mode = "RGBA"
s.view_settings.view_transform = "Standard"
s.view_settings.look = "None"
s.world.color = (0.12, 0.12, 0.12)


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


floor = mat("Slate", (0.11, 0.16, 0.24))
white = mat("Ivory", (0.85, 0.87, 0.89))
orange = mat("Amber", (0.85, 0.44, 0.16))
blue = mat("Cerulean", (0.22, 0.48, 0.72))
# Ground plane extends in Z
bpy.ops.mesh.primitive_plane_add(size=200, location=(0, -1.1, -10))
ground = bpy.context.object
ground.rotation_euler = (math.pi / 2, 0, 0)
# A horizontal floor is XZ with normal +Y
ground.rotation_euler = (math.pi / 2, 0, 0)
ground.data.materials.append(floor)
objects = []
for name, loc, scale, material in [
    ("Near cube", (-0.65, -0.05, -3), (0.35, 0.35, 0.35), orange),
    ("Zero plane cube", (0.65, 0.03, -5), (0.6, 0.6, 0.6), white),
    ("Far cube", (-1.05, 0.35, -8), (0.95, 0.95, 0.95), blue),
]:
    bpy.ops.mesh.primitive_cube_add(size=1, location=loc)
    o = bpy.context.object
    o.name = name
    o.scale = scale
    o.rotation_euler = (0.12, 0.2, 0.08)
    o.data.materials.append(material)
    bev = o.modifiers.new("Soft edges", "BEVEL")
    bev.width = 0.06
    bev.segments = 3
    o.modifiers.new("Weighted normals", "WEIGHTED_NORMAL")
    objects.append(o)
bpy.ops.mesh.primitive_torus_add(
    major_radius=0.55, minor_radius=0.13, location=(0.15, 0.4, -5), rotation=(0, 0, 0)
)
bpy.context.object.data.materials.append(white)
for x, y, z, power, size in [(-3, 5, 1, 450, 5), (4, 2, -2, 320, 4)]:
    bpy.ops.object.light_add(type="AREA", location=(x, y, z))
    light = bpy.context.object
    light.data.energy = power
    light.data.shape = "DISK"
    light.data.size = size
    light.rotation_euler = (Vector((0, 0, -5)) - light.location).to_track_quat("-Z", "Y").to_euler()
cameras = []
for side, x in [("left", -0.065), ("right", 0.065)]:
    bpy.ops.object.camera_add(location=(x, 0, 0))
    c = bpy.context.object
    c.name = "Stereo " + side
    c.data.lens = 50
    c.data.sensor_width = 36
    c.data.sensor_fit = "HORIZONTAL"
    c.data.shift_x = -x * 50 / (5 * 36)
    cameras.append(c)
bpy.context.view_layer.update()
rows = []
for o in objects:
    left = world_to_camera_view(s, cameras[0], o.location)
    right = world_to_camera_view(s, cameras[1], o.location)
    expected = (640 * 50 / 36) * 0.13 * (1 / (-o.location.z) - 1 / 5)
    actual = (left.x - right.x) * 640
    assert abs(actual - expected) < 0.0001, (actual, expected)
    rows.append(
        {
            "object": o.name,
            "depthMeters": -o.location.z,
            "leftCenterX": left.x * 640,
            "rightCenterX": right.x * 640,
            "disparityPixels": actual,
            "formulaPixels": expected,
        }
    )
record = {
    "scope": "Original native Blender two camera render no AI inference",
    "blender": bpy.app.version_string,
    "width": 640,
    "height": 400,
    "baselineMeters": 0.13,
    "lensMM": 50,
    "sensorWidthMM": 36,
    "zeroParallaxMeters": 5,
    "cameraMode": "Parallel optical axes with opposing lens shifts",
    "formula": "width*lens/sensorWidth*baseline*(1/depth-1/zeroPlane)",
    "landmarks": rows,
    "meaning": "Projected object origins from Blender camera matrices not image matched disparity or comfort certification",
}
(out / "anaglyph-camera-study.json").write_text(json.dumps(record, indent=2) + "\n")
with (out / "anaglyph-camera-landmarks.csv").open("w") as f:
    w = csv.DictWriter(f, fieldnames=rows[0].keys(), lineterminator="\n")
    w.writeheader()
    w.writerows(rows)
for side, c in zip(["left", "right"], cameras):
    s.camera = c
    s.render.filepath = str(out / f"anaglyph-{side}.png")
    bpy.ops.render.render(write_still=True)
s.camera = cameras[0]
bpy.ops.wm.save_as_mainfile(filepath=str(out / "anaglyph-original-scene.blend"))
print(json.dumps(record))