import numpy as np
from matplotlib import colormaps
import rerun as rr
from rerun.components import Material
from scipy.spatial import transform


def color_fn(x, cmap="tab10"):
    return colormaps[cmap](x % colormaps[cmap].N)


def ccd_log_sample(
    root_name: str,
    traj: np.ndarray,
):
    
    traj = traj[0]
    num_cameras = traj.shape[0]

    rr.log(root_name, rr.ViewCoordinates.RIGHT_HAND_Y_DOWN, timeless=True)

    rr.log(
        f"{root_name}/trajectory/points",
        rr.Points3D(traj[:, :3]),
        timeless=True,
    )

    rr.log(
        f"{root_name}/trajectory/line",
        rr.LineStrips3D(
            np.stack((traj[:, :3][:-1], traj[:, :3][1:]), axis=1),
            colors=[(1.0, 0.0, 1.0, 1.0)],  # Purple color
        ),
        timeless=True,
    )
    

    for k in range(num_cameras):
        rr.set_time_sequence("frame_idx", k)

        translation = traj[k][:3]

        fx = 955.02  # Focal length in X
        fy = 955.02  # Focal length in Y (same as fx for 1:1 aspect ratio)
        cx = 256     # Principal point X (center of 512x512 image)
        cy = 256     # Principal point Y (center of 512x512 image)
        K = np.array([
            [fx, 0, cx],
            [0, fy, cy],
            [0, 0, 1]
        ])

        rr.log(
            f"{root_name}/camera/image",
            rr.Pinhole(
                image_from_camera=K,
                width=K[0, -1] * 2,
                height=K[1, -1] * 2,
            ),
        )

        width = K[0, -1] * 2 
        height = K[1, -1] * 2
        fov_x = 2 * np.arctan(width / (2 * K[0, 0]))
        fov_y = 2 * np.arctan(height / (2 * K[1, 1]))

        horizontal_angle = np.arctan(traj[k][3] * np.tan(fov_x / 2))
        vertical_angle = np.arctan(traj[k][4] * np.tan(fov_y / 2))

        direction = -translation
        direction /= np.linalg.norm(direction) 

        up = np.array([0, 1, 0])

        right = np.cross(up, direction)
        right /= np.linalg.norm(right)
        up = np.cross(direction, right)

        rotation_matrix = np.vstack([right, up, direction]).T

        rotation_x = transform.Rotation.from_rotvec(vertical_angle * np.array([1, 0, 0])) 
        rotation_y = transform.Rotation.from_rotvec(-horizontal_angle * np.array([0, 1, 0]))

        rotation_combined = rotation_y * transform.Rotation.from_matrix(rotation_matrix) * rotation_x

        rotation_q = rotation_combined.as_quat()

        rr.log(
            f"{root_name}/camera",
            rr.Transform3D(
                translation=translation,
                rotation=rr.Quaternion(xyzw=rotation_q),
            ),
        )

        rr.set_time_sequence("image", k)

        rr.log(
            f"{root_name}/char_traj/points",
            rr.Points3D([[0, 0, 0]], colors=[(1.0, 0.0, 0.0, 1.0)]),  
            timeless=True,
        )