from pathlib import Path
import numpy as np
from custom.replay_utils import create_replay_context, replay_sequence, close_context
import tensorflow_datasets as tfds
import json
from tqdm import tqdm
from scipy.spatial.transform import Rotation as R

camera_pose_np = np.array([[ 0.12289379, -0.25066368,  0.96024207, -0.12524252],
                        [-0.9860285,  -0.14047504,  0.08952408, -0.41773149],
                        [ 0.1124496,  -0.957828,   -0.26442504,  0.41110082],
                        [ 0.,          0.,          0.,          1.        ]]) # 4x4, float64
joint_positions_np = np.random.randn(10, 8) # np.load("droid/proprio_states.npy") # N x 8, float64
# np.array([[0.07692957, -0.62007982, -0.14516954, -2.71003532,  0.07004238,  2.26837683, -0.2618871,   0.0]])
intrinsics_np = np.array([[525.31878662,   0.,         648.12060547],
                        [  0.,         525.31878662, 374.60479736],
                        [  0.,           0.,           1.        ]]) # 3x3, float64

# create the replay context
ctx = create_replay_context(headless=True, scene=1)
try:
    replay_sequence(
        ctx=ctx,
        camera_pose=camera_pose_np,
        intrinsics=intrinsics_np,
        joint_positions=joint_positions_np,
        output_path=Path("runs/debug/replay.mp4"),
    )
except Exception as e:
    raise e
finally:
    close_context(ctx)