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(subject: Dict) -> Dict[str, np.ndarray]:
    position = vector3_to_numpy(subject['position'])
    size = vector3_to_numpy(subject['size'])

    return {
        'center': position,
        'half_size': size / 2
    }


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:
        if not subjects:
            return

        centers = []
        half_sizes = []
        colors = []
        labels = []

        for subject in subjects:
            try:
                box_params = create_subject_box(subject)
                centers.append(box_params['center'])
                half_sizes.append(box_params['half_size'])
                colors.append([0.8, 0.2, 0.2, 1.0])
                labels.append(subject.get('objectClass', 'Unknown'))
            except Exception as e:
                print(f"Error creating box parameters: {str(e)}")
                continue

        if centers:
            rr.log(
                "world/subjects",
                rr.Boxes3D(
                    centers=np.array(centers),
                    half_sizes=np.array(half_sizes),
                    colors=np.array(colors),
                    labels=labels,
                    fill_mode="solid"
                ),
                timeless=True
            )

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

        try:
            camera_positions = np.array([
                vector3_to_numpy(frame['position']) for frame in camera_frames
            ])
            rr.log(
                "world/camera_trajectory",
                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:
        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)}")