import torch
import pypose as pp
from pfp.common.visualization import RerunViewer as RV
from pfp.common.se3_utils import pfp_to_state5p_th, state5p_to_pfp_th, pfp_to_pose_th

B = 2
T = 2
grippers = torch.ones(B, T, 1)
poses = pp.randn_SE3(B, T).matrix()
state_pfp = torch.cat([poses[..., :3, 3], poses[..., :3, 0], poses[..., :3, 1], grippers], dim=-1)
state_5p = pfp_to_state5p_th(state_pfp)
state_5p_pcds = state_5p[:, :, :15].reshape(B, T, 5, 3)

noise = torch.randn_like(state_5p) * 0.01
noise[..., 15] = 0
noisy_state_5p = state_5p + noise
noisy_5p_pcds = noisy_state_5p[:, :, :15].reshape(B, T, 5, 3)
estimated_state_pfp = state5p_to_pfp_th(noisy_state_5p)
estimated_poses, gripper = pfp_to_pose_th(estimated_state_pfp)

RV("state_5p")

for i in range(B):
    for j in range(T):
        RV.add_axis(f"state_pfp_{i}_{j}", poses[i, j].numpy())
        RV.add_axis(f"estimated_{i}_{j}", estimated_poses[i, j].numpy())
        RV.add_np_pointcloud(f"state_5p_{i}_{j}", state_5p_pcds[i, j].numpy(), radii=0.005)
        RV.add_np_pointcloud(f"noisy_5p_{i}_{j}", noisy_5p_pcds[i, j].numpy(), radii=0.005)


print("done")