import rerun as rr
import numpy as np
from typing import Dict, Any, List
from utils.geometry import vector3_to_numpy, euler_to_quaternion


def create_subject_box(frame: Dict[str, Any], dimensions: Dict[str, float]) -> Dict[str, np.ndarray]:
    """
    Given a single frame (position, rotation) and subject's dimensions,
    return the center and half_size for a box representation.
    """
    position = vector3_to_numpy(frame['position'])
    half_size = np.array([
        dimensions['width'] / 2.0,
        dimensions['height'] / 2.0,
        dimensions['depth'] / 2.0
    ], dtype=np.float32)

    return {
        'center': position,
        'half_size': half_size
    }


class SimulationLogger:
    def __init__(self):
        rr.init("camera_simulation")
        rr.log("world", rr.ViewCoordinates.RIGHT_HAND_Y_UP, timeless=True)

        self.K = np.array([
            [500, 0, 960],
            [0, 500, 540],
            [0, 0, 1]
        ])

    def log_metadata(self, instructions: List[Dict[str, Any]]) -> None:
        if not instructions:
            return

        rr.log(
            "metadata/instructions",
            rr.TextDocument(
                "\n".join([
                    f"Instruction {i+1}:\n"
                    f"  Movement: {inst.get('cameraMovement', 'N/A')}\n"
                    f"  Easing: {inst.get('movementEasing', 'N/A')}\n"
                    f"  Frames: {inst.get('frameCount', 'N/A')}\n"
                    f"  Camera Angle: {inst.get('initialCameraAngle', 'N/A')}\n"
                    f"  Shot Type: {inst.get('initialShotType', 'N/A')}\n"
                    f"  Subject Index: {inst.get('subjectIndex', 'N/A')}"
                    for i, inst in enumerate(instructions)
                ])
            ),
            timeless=True
        )

    def log_subjects(self, subjects: List[Dict[str, Any]]) -> None:
        """
        Handles dynamic subjects:
        - Each subject has multiple frames (position, rotation).
        - Subject dimensions (width, height, depth) remain constant across frames.
        """
        if not subjects:
            return

        for subject_idx, subject in enumerate(subjects):
            dimensions = subject.get("dimensions", {})
            frames = subject.get("frames", [])
            label = subject.get('objectClass', f"Subject_{subject_idx}")

            for frame_idx, frame in enumerate(frames):
                rr.set_time_sequence("frame_idx", frame_idx)

                try:
                    box_params = create_subject_box(frame, dimensions)
                    center = box_params['center']
                    half_size = box_params['half_size']

                    # Convert Euler angles to quaternion for rotation
                    rotation_q = euler_to_quaternion(frame['rotation'])

                    # 1) Log a transform for this subject so the box is placed at 'center' with correct rotation
                    rr.log(
                        f"world/subjects/subject_{subject_idx}",
                        rr.Transform3D(
                            translation=center,
                            rotation=rr.Quaternion(xyzw=rotation_q)
                        )
                    )

                    # 2) Log a box at the origin of the subject's local transform
                    rr.log(
                        f"world/subjects/subject_{subject_idx}/box",
                        rr.Boxes3D(
                            centers=np.array([[0.0, 0.0, 0.0]], dtype=np.float32),
                            half_sizes=np.array([half_size], dtype=np.float32),
                            colors=np.array([[0.8, 0.2, 0.2, 1.0]], dtype=np.float32),
                            labels=[label],
                            show_labels=False,
                            fill_mode="solid"
                        )
                    )

                except Exception as e:
                    print(f"Error creating box parameters for subject {subject_idx}, frame {frame_idx}: {str(e)}")

    def log_subject_trajectories(self, subjects: List[Dict[str, Any]]) -> None:
        """
        For each subject, collect all frame centers and log them as:
         - A set of 3D points (subject center positions).
         - A line strip connecting them (showing the subject's center path).
        This is done timelessly so the full trajectory is always visible.
        """
        if not subjects:
            return

        for subject_idx, subject in enumerate(subjects):
            dimensions = subject.get("dimensions", {})
            frames = subject.get("frames", [])
            label = subject.get('objectClass', f"Subject_{subject_idx}")

            if not frames:
                continue

            # Gather center positions
            center_positions = []
            for frame in frames:
                pos = vector3_to_numpy(frame['position'])
                center_positions.append(pos)

            center_positions = np.array(center_positions, dtype=np.float32)

            # Log the points
            rr.log(
                f"world/subjects/subject_{subject_idx}/center_trajectory_points",
                rr.Points3D(
                    center_positions,
                    colors=np.full((len(center_positions), 4), [0.8, 0.6, 0.2, 1.0])
                ),
                timeless=True
            )

            # Log a line strip
            if len(center_positions) > 1:
                lines = np.stack([center_positions[:-1], center_positions[1:]], axis=1)
                rr.log(
                    f"world/subjects/subject_{subject_idx}/center_trajectory_line",
                    rr.LineStrips3D(
                        lines,
                        colors=[(0.8, 0.6, 0.2, 1.0)],
                    ),
                    timeless=True
                )

    def log_camera_trajectory(self, camera_frames: List[Dict[str, Any]]) -> None:
        """
        Logs the entire camera trajectory as a set of 3D points and a connecting line strip.
        This is typically done timelessly to show the full path at once.
        """
        if not camera_frames:
            return

        try:
            camera_positions = np.array([
                vector3_to_numpy(frame['position']) for frame in camera_frames
            ], dtype=np.float32)

            rr.log(
                "world/camera_trajectory_points",
                rr.Points3D(
                    camera_positions,
                    colors=np.full((len(camera_positions), 4), [0.0, 0.8, 0.8, 1.0])
                ),
                timeless=True
            )

            if len(camera_positions) > 1:
                lines = np.stack([camera_positions[:-1], camera_positions[1:]], axis=1)
                rr.log(
                    "world/camera_trajectory_line",
                    rr.LineStrips3D(
                        lines,
                        colors=[(0.0, 0.8, 0.8, 1.0)]
                    ),
                    timeless=True
                )

        except Exception as e:
            print(f"Error logging camera trajectory: {str(e)}")

    def log_camera_frames(self, camera_frames: List[Dict[str, Any]]) -> None:
        """
        Logs each camera frame at a given time sequence, so you can watch it update over time.
        """
        if not camera_frames:
            return

        for frame_idx, camera_frame in enumerate(camera_frames):
            try:
                rr.set_time_sequence("frame_idx", frame_idx)

                position = vector3_to_numpy(camera_frame['position'])
                rotation_q = euler_to_quaternion(camera_frame['angle'])

                rr.log(
                    "world/camera",
                    rr.Transform3D(
                        translation=position,
                        rotation=rr.Quaternion(xyzw=rotation_q)
                    )
                )

                rr.log(
                    "world/camera/image",
                    rr.Pinhole(
                        image_from_camera=self.K,
                        width=1920,
                        height=1080
                    )
                )

            except Exception as e:
                print(f"Error logging camera frame {frame_idx}: {str(e)}")

    def log_helper_keyframes(self, helper_keyframes: List[Dict[str, Any]]) -> None:
        if not helper_keyframes:
            return

        helper_positions = np.array([
            vector3_to_numpy(frame['position']) for frame in helper_keyframes
        ])
        rr.log(
            "world/helper_keyframes",
            rr.Points3D(
                helper_positions,
                radii=np.full(len(helper_positions), 0.03),
                colors=np.full((len(helper_positions), 4), [1.0, 1.0, 0.0, 1.0]),
            ),
            timeless=True
        )

        for keyframe_idx, helper_keyframe in enumerate(helper_keyframes):
            try:
                position = vector3_to_numpy(helper_keyframe['position'])
                rotation_q = euler_to_quaternion(helper_keyframe['angle'])

                rr.log(
                    f"world/helper_camera_{keyframe_idx}",
                    rr.Transform3D(
                        translation=position,
                        rotation=rr.Quaternion(xyzw=rotation_q),
                        scale=(.5, .5, .5)
                    ),
                    timeless=True
                )

                rr.log(
                    f"world/helper_camera_{keyframe_idx}/image",
                    rr.Pinhole(
                        image_from_camera=self.K,
                        width=1920,
                        height=1080,
                    )
                )

            except Exception as e:
                print(f"Error logging helper keyframe {keyframe_idx}: {str(e)}")