"""Original Clout native studies for Blender 5 Eevee
Run: blender --background --python render-world-vertex-study.py -- /new/output
No AI inference or Mundo reconstruction, renderer or MuJoCo execution
"""

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

import bpy
from mathutils import Vector

out = Path(sys.argv[sys.argv.index("--") + 1])
out.mkdir(parents=True, exist_ok=False)


def setup(width=960, height=540):
    bpy.ops.wm.read_factory_settings(use_empty=True)
    s = bpy.context.scene
    s.render.engine = "BLENDER_EEVEE"
    s.render.resolution_x = width
    s.render.resolution_y = height
    s.render.resolution_percentage = 100
    s.render.image_settings.file_format = "PNG"
    s.render.fps = 24
    s.world = bpy.data.worlds.new("Studio")
    s.world.color = (0.10, 0.10, 0.10)
    s.view_settings.view_transform = "AgX"
    return s


def camera(s, loc, target, scale):
    bpy.ops.object.camera_add(location=loc)
    c = bpy.context.object
    c.rotation_euler = (Vector(target) - c.location).to_track_quat("-Z", "Y").to_euler()
    c.data.type = "ORTHO"
    c.data.ortho_scale = scale
    s.camera = c
    return c


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


def lights():
    for loc, energy in [((1, -4, 7), 1300), ((-4, 2, 5), 700)]:
        bpy.ops.object.light_add(type="AREA", location=loc)
        o = bpy.context.object
        o.data.energy = energy
        o.data.size = 5
        o.rotation_euler = (-o.location).to_track_quat("-Z", "Y").to_euler()


records = {
    "kind": "Original native color-attribute and rigid-body studies, no AI inference",
    "blender": bpy.app.version_string,
}
for n in [2, 32]:
    s = setup(900, 720)
    verts = [(-1.5 + 3 * i / n, -1 + 2 * j / n, 0) for j in range(n + 1) for i in range(n + 1)]
    faces = [
        (j * (n + 1) + i, j * (n + 1) + i + 1, (j + 1) * (n + 1) + i + 1, (j + 1) * (n + 1) + i)
        for j in range(n)
        for i in range(n)
    ]
    mesh = bpy.data.meshes.new("Explicit paint grid")
    mesh.from_pydata(verts, [], faces)
    mesh.update()
    obj = bpy.data.objects.new("Vertex color surface", mesh)
    s.collection.objects.link(obj)
    attr = mesh.color_attributes.new(name="OriginalPaint", type="FLOAT_COLOR", domain="POINT")
    for i, (x, y, z) in enumerate(verts):
        d = math.sqrt((x / 0.70) ** 2 + (y / 0.50) ** 2)
        a = max(0, min(1, (1 - d) / 0.25))
        attr.data[i].color = (
            0.025 * (1 - a) + 0.8 * a,
            0.23 * (1 - a) + 0.25 * a,
            0.3 * (1 - a) + 0.035 * a,
            1,
        )
    mat = bpy.data.materials.new("Point color material")
    mat.use_nodes = True
    nodes = mat.node_tree.nodes
    nodes.clear()
    col = nodes.new("ShaderNodeVertexColor")
    col.layer_name = "OriginalPaint"
    em = nodes.new("ShaderNodeEmission")
    output = nodes.new("ShaderNodeOutputMaterial")
    mat.node_tree.links.new(col.outputs["Color"], em.inputs["Color"])
    mat.node_tree.links.new(em.outputs[0], output.inputs["Surface"])
    obj.data.materials.append(mat)
    camera(s, (0, 0, 5), (0, 0, 0), 3.8)
    name = "vertex-paint-coarse" if n == 2 else "vertex-paint-dense"
    records[name] = {
        "subdivisions": n,
        "vertices": len(verts),
        "faces": len(faces),
        "UVLayers": len(mesh.uv_layers),
        "attributeDomain": attr.domain,
        "attributeType": attr.data_type,
        "colorHash": hashlib.sha256(
            json.dumps([list(d.color) for d in attr.data]).encode()
        ).hexdigest(),
        "formula": "a=clamp((1-sqrt((x/0.70)^2+(y/0.50)^2))/0.25,0,1), interpolate teal to orange at vertices",
    }
    s.render.filepath = str(out / (name + ".png"))
    bpy.ops.wm.save_as_mainfile(filepath=str(out / (name + ".blend")))
    bpy.ops.render.render(write_still=True)

s = setup()
s.frame_start = 1
s.frame_end = 48


def box(name, loc, dimensions, color, dynamic=False):
    bpy.ops.mesh.primitive_cube_add(size=2, location=loc)
    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(material(name, color))
    bpy.ops.rigidbody.object_add()
    o.rigid_body.type = "ACTIVE" if dynamic else "PASSIVE"
    o.rigid_body.collision_shape = "BOX"
    o.rigid_body.friction = 0.6
    o.rigid_body.restitution = 0.1
    o.rigid_body.use_margin = True
    o.rigid_body.collision_margin = 0.001
    return o


floor = box("Floor", (0, 0, -0.1), (8, 6, 0.2), (0.10, 0.15, 0.20))
pedestal = box("Pedestal", (0, 0, 0.3), (1.4, 1.4, 0.6), (0.04, 0.42, 0.38))
parcel = box("Parcel", (0.1, 0, 2.4), (0.6, 0.6, 0.6), (0.85, 0.25, 0.06), True)
parcel.rigid_body.mass = 1
s.rigidbody_world.substeps_per_frame = 5
s.rigidbody_world.solver_iterations = 20
s.rigidbody_world.point_cache.frame_end = 48
cam = camera(s, (6, -8, 6), (0, 0, 0.8), 9)
lights()
textmat = bpy.data.materials.new("HUD")
textmat.use_nodes = True
nodes = textmat.node_tree.nodes
nodes.clear()
em = nodes.new("ShaderNodeEmission")
em.inputs["Color"].default_value = (0.8, 0.9, 1, 1)
output = nodes.new("ShaderNodeOutputMaterial")
textmat.node_tree.links.new(em.outputs[0], output.inputs["Surface"])
for text, y in [
    ("Original Blender rigid body proxy", 2.12),
    ("Orange parcel and teal pedestal | No Mundo inference", -2.22),
]:
    bpy.ops.object.text_add()
    o = bpy.context.object
    o.data.body = text
    o.data.size = 0.19
    o.data.materials.append(textmat)
    o.parent = cam
    o.location = (-4.15, y, -5)
    o.rotation_euler = (0, 0, 0)
positions = []
matrices = []
for f in range(1, 49):
    s.frame_set(f)
    bpy.context.view_layer.update()
    m = parcel.evaluated_get(bpy.context.evaluated_depsgraph_get()).matrix_world.copy()
    matrices.append(m)
    loc, q, scale = m.decompose()
    positions.append(
        {
            "frame": f,
            "time_seconds": (f - 1) / 24,
            "x": loc.x,
            "y": loc.y,
            "z": loc.z,
            "rotation": list(q),
        }
    )
bpy.context.view_layer.objects.active = parcel
parcel.select_set(True)
bpy.ops.rigidbody.object_remove()
parcel.rotation_mode = "QUATERNION"
for f, m in enumerate(matrices, 1):
    loc, q, scale = m.decompose()
    parcel.location = loc
    parcel.rotation_quaternion = q
    parcel.keyframe_insert("location", frame=f)
    parcel.keyframe_insert("rotation_quaternion", frame=f)
for fc in (
    parcel.animation_data.action.layers[0]
    .strips[0]
    .channelbag(parcel.animation_data.action_slot)
    .fcurves
):
    for point in fc.keyframe_points:
        point.interpolation = "LINEAR"
s.frame_set(1)
views = [(6, -8, 6), (-6, -6, 4), (2, 8, 5)]
for i, loc in enumerate(views):
    cam.location = loc
    cam.rotation_euler = (Vector((0, 0, 0.8)) - cam.location).to_track_quat("-Z", "Y").to_euler()
    s.render.filepath = str(out / f"proxy-context-{i + 1}.png")
    bpy.ops.render.render(write_still=True)
cam.location = views[0]
cam.rotation_euler = (Vector((0, 0, 0.8)) - cam.location).to_track_quat("-Z", "Y").to_euler()
s.frame_set(1)
bpy.ops.wm.save_as_mainfile(filepath=str(out / "world-proxy-study.blend"))
(out / "frames").mkdir()
s.render.filepath = str(out / "frames" / "proxy-")
bpy.ops.render.render(animation=True)
with (out / "world-proxy-positions.csv").open("w") as file:
    writer = csv.writer(file)
    writer.writerow(["frame", "time_seconds", "x", "y", "z", "qw", "qx", "qy", "qz"])
    for r in positions:
        writer.writerow([r["frame"], r["time_seconds"], r["x"], r["y"], r["z"], *r["rotation"]])
records["worldProxy"] = {
    "backend": "Blender Bullet rigid body, baked to parcel location and quaternion keys",
    "frames": 48,
    "fps": 24,
    "width": 960,
    "height": 540,
    "duration": 2,
    "parcelDimensions": [0.6, 0.6, 0.6],
    "pedestalDimensions": [1.4, 1.4, 0.6],
    "pedestalTop": 0.6,
    "gravity": list(s.gravity),
    "startZ": positions[0]["z"],
    "endZ": positions[-1]["z"],
    "minZ": min(r["z"] for r in positions),
    "cameraViews": views,
    "positions": positions,
    "meaning": "Original authored scene, no Mundo reconstruction or video model inference and no MuJoCo execution",
}
(out / "world-vertex-measurements.json").write_text(json.dumps(records, indent=2) + "\n")
