import torch
import kaolin as kal
from kaolin.render.camera import Camera
from kaolin.render.easy_render import render_mesh, default_lighting
from PIL import Image
import math

# デバイスの設定（CUDAが利用可能な場合はGPUを使用）
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# GLBファイルのパス
model_path = 'sample.glb'

# メッシュの読み込み
mesh = kal.io.gltf.import_mesh(model_path)
mesh = mesh.to(device)

# カメラの設定関数
def create_camera(elev, azim, width=1024, height=1024):
    # カメラ位置の計算
    radius = 3.0  # カメラとモデルの距離
    azim_rad = math.radians(azim)
    elev_rad = math.radians(elev)
    eye = torch.tensor([
        radius * math.cos(elev_rad) * math.sin(azim_rad),
        radius * math.cos(elev_rad) * math.cos(azim_rad),
        radius * math.sin(elev_rad)
    ], device=device)
    at = torch.tensor([0.0, 0.0, 0.0], device=device)
    up = torch.tensor([0.0, 1, 0])  # デバイス指定なしでY軸を上方向とする

    # カメラの生成
    camera = Camera.from_args(
        eye=eye,
        at=at,
        up=up,  # `up` ベクトルはデバイス指定なしで渡す
        fov=45.0 * math.pi / 180.0,  # ラジアンに変換
        width=width,
        height=height,
        near=0.1,
        far=10.0,
        dtype=torch.float32,
        device=device
    )
    return camera

# カメラ位置の設定
camera_azims = [-180, -135, -90, -45, 0, 45, 90, 135]
camera_poses = [(0, azim) for azim in camera_azims]
camera_poses.append((30, 0))   # 上からの視点1
camera_poses.append((30, 180)) # 上からの視点2

# ライティングの設定
lighting = default_lighting()
lighting = lighting.to(device)

# 全カメラ位置からレンダリング
for idx, (elev, azim) in enumerate(camera_poses):
    # 高解像度のカメラを生成
    camera = create_camera(elev, azim, width=1024, height=1024)

    # メッシュのレンダリング
    render_dict = render_mesh(
        camera=camera,
        mesh=mesh,
        lighting=lighting
    )

    # レンダリング結果の取得
    rendered_image = render_dict['render'].clamp(0,1)
    image = rendered_image[0].cpu().numpy()
    image = (image * 255).astype('uint8')
    image = Image.fromarray(image)

    # 画像を保存
    output_path = f'rendered_knight_view_{idx}.png'
    image.save(output_path)
    print(f'カメラ位置 {idx} (elev={elev}, azim={azim}) のレンダリングを {output_path} に保存しました。')

print('すべてのレンダリングが完了しました。')