"""Original Clout geometry and UV studies, Blender 5 Eevee, no AI inference."""

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():
    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 = 900
    scene.render.resolution_y = 720
    scene.render.resolution_percentage = 100
    scene.render.image_settings.file_format = "PNG"
    scene.world.color = (0.15, 0.15, 0.15)
    scene.view_settings.view_transform = "AgX"
    bpy.ops.object.camera_add(location=(4, -6, 4.5))
    cam = bpy.context.object
    cam.rotation_euler = (Vector((0, 0, 0)) - cam.location).to_track_quat("-Z", "Y").to_euler()
    cam.data.type = "ORTHO"
    cam.data.ortho_scale = 4.5
    scene.camera = cam
    for loc, power, size in [((0, -3, 5), 1000, 4), ((3, 2, 4), 500, 3)]:
        bpy.ops.object.light_add(type="AREA", location=loc)
        bpy.context.object.data.energy = power
        bpy.context.object.data.shape = "DISK"
        bpy.context.object.data.size = size
        bpy.context.object.rotation_euler = (
            (Vector((0, 0, 0)) - bpy.context.object.location).to_track_quat("-Z", "Y").to_euler()
        )
    scene.render.film_transparent = False
    return scene


def finish(scene, name):
    scene.render.filepath = str(OUT / f"{name}.png")
    bpy.ops.wm.save_as_mainfile(filepath=str(OUT / f"{name}.blend"))
    bpy.ops.render.render(write_still=True)


records = {
    "kind": "Original native Blender study, no model inference",
    "blender": bpy.app.version_string,
}
for displaced in [False, True]:
    scene = setup()
    n = 160
    verts = []
    for j in range(n + 1):
        for i in range(n + 1):
            x = -1.4 + 2.8 * i / n
            y = -1.4 + 2.8 * j / n
            z = 0.13 * math.sin(10 * x) * math.sin(10 * y) if displaced else 0
            verts.append((x, y, z))
    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("Analytic grid")
    mesh.from_pydata(verts, [], faces)
    mesh.update()
    obj = bpy.data.objects.new("Surface", mesh)
    scene.collection.objects.link(obj)
    for poly in mesh.polygons:
        poly.use_smooth = True
    mat = bpy.data.materials.new("Teal surface")
    mat.diffuse_color = (0.05, 0.45, 0.38, 1)
    mat.use_nodes = True
    nodes = mat.node_tree.nodes
    links = mat.node_tree.links
    bs = nodes.get("Principled BSDF")
    bs.inputs["Base Color"].default_value = (0.05, 0.45, 0.38, 1)
    bs.inputs["Roughness"].default_value = 0.45
    if not displaced:
        coord = nodes.new("ShaderNodeTexCoord")
        sep = nodes.new("ShaderNodeSeparateXYZ")
        links.new(
            coord.outputs["Geometry"] if "Geometry" in coord.outputs else coord.outputs["Object"],
            sep.inputs[0],
        )
        sins = []
        for axis in ["X", "Y"]:
            scale = nodes.new("ShaderNodeMath")
            scale.operation = "MULTIPLY"
            scale.inputs[1].default_value = 10
            links.new(sep.outputs[axis], scale.inputs[0])
            sin = nodes.new("ShaderNodeMath")
            sin.operation = "SINE"
            links.new(scale.outputs[0], sin.inputs[0])
            sins.append(sin)
        mul = nodes.new("ShaderNodeMath")
        mul.operation = "MULTIPLY"
        links.new(sins[0].outputs[0], mul.inputs[0])
        links.new(sins[1].outputs[0], mul.inputs[1])
        bump = nodes.new("ShaderNodeBump")
        bump.inputs["Distance"].default_value = 0.13
        links.new(mul.outputs[0], bump.inputs["Height"])
        links.new(bump.outputs["Normal"], bs.inputs["Normal"])
    obj.data.materials.append(mat)
    name = "displacement-geometry" if displaced else "displacement-bump"
    records[name] = {
        "vertices": len(verts),
        "z_min": min(v[2] for v in verts),
        "z_max": max(v[2] for v in verts),
        "height_function": "0.13 sin(10x) sin(10y)",
        "bump_distance": None if displaced else 0.13,
    }
    finish(scene, name)

for stretched in [False, True]:
    scene = setup()
    bpy.ops.mesh.primitive_cube_add(size=2)
    obj = bpy.context.object
    obj.name = "Rectangular package"
    obj.dimensions = (2.4, 1.2, 0.8)
    bpy.ops.object.transform_apply(location=False, rotation=False, scale=True)
    for edge in obj.data.edges:
        edge.use_seam = True
    bpy.ops.object.mode_set(mode="EDIT")
    bpy.ops.mesh.select_all(action="SELECT")
    bpy.ops.uv.unwrap(method="ANGLE_BASED", margin=0.035)
    bpy.ops.uv.average_islands_scale()
    bpy.ops.uv.pack_islands(rotate=False, margin=0.04)
    bpy.ops.object.mode_set(mode="OBJECT")
    uv = obj.data.uv_layers.active.data
    if stretched:
        for face in obj.data.polygons:
            for idx, xy in zip(face.loop_indices, [(0, 0), (1, 0), (1, 1), (0, 1)]):
                uv[idx].uv = xy
    mat = bpy.data.materials.new("UV checker")
    mat.use_nodes = True
    nodes = mat.node_tree.nodes
    links = mat.node_tree.links
    tex = nodes.new("ShaderNodeTexCoord")
    check = nodes.new("ShaderNodeTexChecker")
    check.inputs["Scale"].default_value = 12
    check.inputs["Color1"].default_value = (0.04, 0.45, 0.38, 1)
    check.inputs["Color2"].default_value = (0.85, 0.8, 0.64, 1)
    links.new(tex.outputs["UV"], check.inputs["Vector"])
    bs = nodes.get("Principled BSDF")
    links.new(check.outputs["Color"], bs.inputs["Base Color"])
    bs.inputs["Roughness"].default_value = 0.65
    obj.data.materials.append(mat)
    name = "uv-stretched" if stretched else "uv-unwrapped"
    faces = []
    for face in obj.data.polygons:
        pts = [list(uv[i].uv) for i in face.loop_indices]
        area = (
            abs(
                sum(
                    pts[k][0] * pts[(k + 1) % 4][1] - pts[(k + 1) % 4][0] * pts[k][1]
                    for k in range(4)
                )
            )
            / 2
        )
        faces.append({"face": face.index, "world_area": face.area, "uv_area": area, "uv": pts})
    records[name] = {
        "method": "each face forced into a unit square"
        if stretched
        else "ANGLE_BASED unwrap, average island scale, pack",
        "faces": faces,
    }
    finish(scene, name)
(OUT / "texture-study-measurements.json").write_text(json.dumps(records, indent=2) + "\n")
