from pathlib import Path

import bpy
import numpy as np
import math
from mathutils import Matrix, Vector

# Hardcode your NPZ path here. Use raw string or forward slashes to avoid backslash escapes.
NPZ_PATH = Path(r"D:\Users\Adam\Downloads\VRM_JG6Z7WA_clip_0_global_smplx.npz")
CAMERA_NAME = "GenmoCamera"
SPHERE_NAME = "TestSphere"
FPS = 30
ROTATE_LOCAL_Z_DEG = 180.0
ROTATE_LOCAL_Y_DEG = 180.0


def _ensure_camera(name: str) -> bpy.types.Object:
    cam_obj = bpy.data.objects.get(name)
    if cam_obj is None:
        cam_data = bpy.data.cameras.new(name)
        cam_obj = bpy.data.objects.new(name, cam_data)
        bpy.context.collection.objects.link(cam_obj)
    return cam_obj


def _apply_fov(cam_obj: bpy.types.Object, npz_data: np.lib.npyio.NpzFile) -> None:
    cam_data = cam_obj.data
    fov_y_deg = float(npz_data["fov_y_deg"]) if "fov_y_deg" in npz_data else None
    fov_x_deg = float(npz_data["fov_x_deg"]) if "fov_x_deg" in npz_data else None

    if (fov_x_deg is None or fov_y_deg is None) and "K_fullimg" in npz_data:
        K_fullimg = npz_data["K_fullimg"]
        if K_fullimg.ndim == 3:
            K0 = K_fullimg[0]
        else:
            K0 = K_fullimg
        fx = float(K0[0, 0])
        fy = float(K0[1, 1])
        cx = float(K0[0, 2])
        cy = float(K0[1, 2])
        width = 2.0 * cx
        height = 2.0 * cy
        if fov_x_deg is None and fx > 0 and width > 0:
            fov_x_deg = math.degrees(2.0 * math.atan(width / (2.0 * fx)))
        if fov_y_deg is None and fy > 0 and height > 0:
            fov_y_deg = math.degrees(2.0 * math.atan(height / (2.0 * fy)))

    if fov_y_deg is not None:
        cam_data.angle_y = math.radians(float(fov_y_deg))
    elif fov_x_deg is not None:
        cam_data.angle = math.radians(float(fov_x_deg))


def _ensure_sphere(name: str) -> bpy.types.Object:
    obj = bpy.data.objects.get(name)
    if obj is None:
        bpy.ops.mesh.primitive_uv_sphere_add(radius=0.2, location=(0.0, 0.0, 0.0))
        obj = bpy.context.active_object
        obj.name = name
    return obj


def main():
    npz_path = NPZ_PATH
    if not npz_path.is_absolute():
        npz_path = Path(bpy.path.abspath("//")) / npz_path
    data = np.load(str(npz_path))
    if "camera_transform" not in data:
        raise ValueError("NPZ is missing 'camera_transform'")
    camera_transform = data["camera_transform"]  # (F, 4, 4)

    scene = bpy.context.scene
    scene.render.fps = FPS
    scene.frame_start = 1
    scene.frame_end = int(camera_transform.shape[0])

    cam_obj = _ensure_camera(CAMERA_NAME)
    scene.camera = cam_obj
    _ensure_sphere(SPHERE_NAME)
    _apply_fov(cam_obj, data)

    for i, T in enumerate(camera_transform):
        frame = i + 1
        mat = Matrix(T.tolist())
        rot_z = Matrix.Rotation(math.radians(ROTATE_LOCAL_Z_DEG), 3, Vector((0.0, 0.0, 1.0)))
        rot_y = Matrix.Rotation(math.radians(ROTATE_LOCAL_Y_DEG), 3, Vector((0.0, 1.0, 0.0)))
        rot_local = rot_y @ rot_z  # first local Z, then local Y
        new_rot = mat.to_3x3() @ rot_local
        new_mat = new_rot.to_4x4()
        new_mat.translation = mat.to_translation()
        cam_obj.matrix_world = new_mat
        cam_obj.keyframe_insert(data_path="location", frame=frame)
        cam_obj.keyframe_insert(data_path="rotation_euler", frame=frame)


if __name__ == "__main__":
    main()