import numpy as np
import genesis as gs
import cv2
import imageio

gs.init(backend=gs.cpu)

DT = 2e-3

scene = gs.Scene(
    show_viewer=True,
    sim_options = gs.options.SimOptions(
        dt     = DT,
        gravity= (0.0, -9.81, 0.0),  # 重力ベクトル
    ),
    viewer_options=gs.options.ViewerOptions(
        res=(960, 720),
        camera_pos=(0.0, 4.0, 4.0),
        camera_lookat=(0.0, 0.0, 0.0),
        camera_up=(0.0, 1.0, 0.0),
        camera_fov=15,
        max_FPS=60,
    ),
    renderer=gs.renderers.Rasterizer(),
)

plane = scene.add_entity(
   gs.morphs.Plane(
        pos   = (0.0, 0.0, 0.0),
        euler = (-90.0,   # roll: 90°
                 0.0,    # pitch
                 0.0),   # yaw
   )
)

cobotta = scene.add_entity(
    gs.morphs.MJCF(file='./cobotta.xml'),
)

cam = scene.add_camera(
        res=(960, 720),
        pos=(0.0, 4.0, 4.0),
        lookat=(0.0, 0.0, 0),
        up=(0.0, 1.0, 0.0),
        fov=15,
        GUI=False,
)

scene.build()

jnt_names = [
    'joint1',
    'joint2',
    'joint3',
    'joint4',
    'joint5',
    'joint6',
    'joint7',
    'joint8',
]

dofs_idx = [cobotta.get_joint(name).dof_idx_local for name in jnt_names]

# 7. 動画保存設定
rgb, depth, seg, normal = cam.render(depth=True, segmentation=True, normal=True)
width, height = rgb.shape[1], rgb.shape[0]
fps = 30
video = cv2.VideoWriter('cobotta.mp4',
                        cv2.VideoWriter_fourcc(*'mp4v'),
                        fps,
                        (width, height))

# 8. 録画ループ
cam.start_recording()

for i in range(1000):
    cobotta.set_dofs_position(np.array([180/180*np.pi*np.sin(2*np.pi*0.5*i*DT),
					 90/180*np.pi*np.sin(2*np.pi*1.0*i*DT),
					 45/180*np.pi*np.sin(2*np.pi*1.0*i*DT),
					 45/180*np.pi*np.sin(2*np.pi*1.0*i*DT),
					 45/180*np.pi*np.sin(2*np.pi*1.0*i*DT),
					 45/180*np.pi*np.sin(2*np.pi*1.0*i*DT),
				 	 0.01*np.sin(2*np.pi*i*DT), -0.01*np.sin(2*np.pi*i*DT)	# 1/180*np.pi*i
				]), dofs_idx)
    scene.step()
    cam.set_pose(
            pos=(2.0, -1.0, 4.0), # pos=(3.0 * np.sin(i / 60), 3.0 * np.cos(i / 60), 2.5),
            lookat=(0.0, 0.0, 0.1),
    )
    cam.render()  # 録画モードで内部的にフレーム保存

    # フレームを取得して動画に書き込み
    rgb_frame, _, _, _ = cam.render(depth=False, segmentation=False, normal=False)
    frame_bgr = cv2.cvtColor(rgb_frame, cv2.COLOR_RGB2BGR)
    video.write(frame_bgr)

video.release()

input()