"""
Template for converting a custom dataset to MIMIC format.

This example shows the basic structure and steps needed to convert
any autonomous driving dataset to the MIMIC data format.

Steps:
1. Load raw data from your dataset
2. Parse sensor calibration (intrinsics/extrinsics)
3. Create mimic-data-api objects
4. Handle coordinate system transformations
5. Serialize to JSON and save data files

Replace the placeholder functions with your actual data parsing logic.
"""

import json
import numpy as np
from pathlib import Path
from typing import Dict, List, Tuple
import argparse
from scipy.spatial.transform import Rotation, Slerp

from mimic_data_api.scene import Scene, EgoVehicle, Sensors
from mimic_data_api.camera import (
    Camera,
    CameraIntrinsics,
    DistortionParameters,
    ImageData,
)
from mimic_data_api.lidar import (
    Lidar,
    LidarIntrinsics,
    VerticalResolution,
    LidarPointCloudData,
)
from mimic_data_api.actor import Actor, ActorDimensions
from mimic_data_api.trajectory_numpy import WorldTrajectory


# ============================================================================
# PART 1: Data Loading
# ============================================================================

def load_raw_dataset(dataset_path: Path) -> Dict:
    """
    Load your raw dataset files.
    
    Replace this with actual file loading logic for your format.
    Should return a dictionary with:
    - timestamps: frame timestamps
    - sensor_data: image and point cloud paths
    - calibration: intrinsics/extrinsics
    - annotations: actor/annotation data
    """
    print(f"Loading dataset from {dataset_path}")
    
    # TODO: Implement your dataset loading
    # Example structure:
    dataset = {
        "timestamps": [0.0, 0.1, 0.2, 0.3],  # Frame timestamps in seconds
        "camera_timestamps": {},  # Optional: camera ID -> that camera's image timestamps
        "frames": [],  # Per-frame data
        "ego_poses": [],  # Ego trajectory
        "actors": [],  # Annotations
        "calibration": {},  # Sensor calibration
    }
    
    return dataset


# ============================================================================
# PART 2: Coordinate Transformation
# ============================================================================

def transform_dataset_to_mimic(
    source_to_target: np.ndarray,
) -> np.ndarray:
    """
    Apply coordinate system transformation from dataset format to MIMIC.
    
    Your dataset likely uses different coordinate conventions.
    Create the appropriate transformation matrix here.
    
    Example: Waymo (X=forward, Y=right, Z=down) → MIMIC (X=forward, Y=left, Z=up)
    """
    # Define your transformation matrix (4x4 homogeneous)
    # This example is identity (no transformation)
    return source_to_target


# ============================================================================
# PART 3: Sensor Calibration Parsing
# ============================================================================

class SensorCalibration:
    """Parse sensor calibration from your dataset format."""

    def __init__(self, calib_data: Dict):
        self.calib_data = calib_data

    def get_camera_intrinsics(self, camera_id: int) -> CameraIntrinsics:
        """
        Parse camera intrinsics from your dataset.
        
        You need to provide:
        - resolution: [width, height]
        - focal_length: [fx, fy]
        - principal_point: [cx, cy]
        - distortion: radial-tangential or fisheye parameters
        """
        # TODO: Implement camera intrinsics parsing
        
        # Example:
        intrinsics = CameraIntrinsics(
            fx=1000.0,
            fy=1000.0,
            cx=960.0,
            cy=540.0,
            skew=0.0,
            width=1920.0,
            height=1080.0,
        )
        
        return intrinsics

    def get_camera_extrinsics(self, camera_id: int) -> np.ndarray:
        """
        Parse camera extrinsics (camera-to-vehicle transform).
        
        Return a 4x4 transformation matrix from camera frame to vehicle frame.
        Apply coordinate system transformation if needed.
        """
        # TODO: Implement camera extrinsics parsing
        
        # Start with identity (camera at vehicle origin)
        extrinsics = np.eye(4)
        
        # Example: Offset camera forward by 1.5m
        extrinsics[0, 3] = 1.5
        
        # Apply coordinate transformation if needed
        # extrinsics = transform_dataset_to_mimic(extrinsics)
        
        return extrinsics

    def get_lidar_intrinsics(self, lidar_id: int) -> LidarIntrinsics:
        """
        Parse LiDAR intrinsics (field of view, resolution, range).
        
        For rotating LiDARs (e.g., Velodyne):
        - vertical_resolution should include elevation_channel_mapping
        
        For solid-state LiDARs:
        - vertical_resolution can use uniform resolution
        """
        # TODO: Implement LiDAR intrinsics parsing
        
        # Example: Velodyne HDL-64E
        intrinsics = LidarIntrinsics(
            lidar_type="rotating",
            rotation_type="clockwise",
            vertical_fov=0.7,  # ~40 degrees
            horizontal_fov=2 * np.pi,  # 360 degrees
            horizontal_resolution=0.002,  # radians/point (angular step), not points/radian
            vertical_resolution=VerticalResolution(
                # Per-channel elevation mapping for HDL-64E
                elevation_channel_mapping={
                    0: -0.45, 1: -0.37, 2: -0.29, 3: -0.21,
                    # ... (64 channels total - provide all)
                }
            ),
            max_range=120.0,
            min_range=0.5,
        )
        
        return intrinsics

    def get_lidar_extrinsics(self, lidar_id: int) -> np.ndarray:
        """
        Parse LiDAR extrinsics (LiDAR-to-vehicle transform).
        
        Return a 4x4 transformation matrix from LiDAR frame to vehicle frame.
        """
        # TODO: Implement LiDAR extrinsics parsing
        
        extrinsics = np.eye(4)
        
        # Example: LiDAR mounted at vehicle center, 2m above ground
        extrinsics[2, 3] = 2.0
        
        # Apply coordinate transformation if needed
        # extrinsics = transform_dataset_to_mimic(extrinsics)
        
        return extrinsics


# ============================================================================
# PART 4: Create Scene Components
# ============================================================================

def create_ego_trajectory(
    timestamps: np.ndarray,
    poses: np.ndarray,  # N x 4 x 4 matrices
) -> WorldTrajectory:
    """
    Create ego vehicle trajectory from poses.
    
    Args:
        timestamps: Array of frame timestamps (N,)
        poses: Array of vehicle-to-world transformation matrices (N, 4, 4)
    
    Returns:
        WorldTrajectory object
    """
    return WorldTrajectory(
        timestamps=timestamps,
        object2world=poses,
    )


def interpolate_trajectory(
    trajectory: WorldTrajectory,
    query_timestamps: np.ndarray,
) -> WorldTrajectory:
    """Interpolate poses within the source trajectory's timestamp coverage.

    Translation is linear and rotation uses spherical linear interpolation
    (SLERP). Requires scipy. Use source poses directly when they are already
    sampled at the target sensor timestamps.
    """
    source_timestamps = np.asarray(trajectory.timestamps, dtype=float)
    query_timestamps = np.asarray(query_timestamps, dtype=float)
    source_poses = np.asarray(trajectory.object2world, dtype=float)

    source_indices = np.searchsorted(source_timestamps, query_timestamps)
    if np.all(source_indices < len(source_timestamps)) and np.array_equal(
        source_timestamps[source_indices], query_timestamps
    ):
        return WorldTrajectory(
            timestamps=query_timestamps,
            object2world=source_poses[source_indices],
        )

    if len(source_timestamps) < 2 or np.any(np.diff(source_timestamps) <= 0):
        raise ValueError("Interpolation requires at least two increasing source timestamps")
    if np.any(query_timestamps < source_timestamps[0]) or np.any(
        query_timestamps > source_timestamps[-1]
    ):
        raise ValueError("Cannot interpolate outside the source trajectory timestamp range")

    translations = np.column_stack(
        [
            np.interp(query_timestamps, source_timestamps, source_poses[:, axis, 3])
            for axis in range(3)
        ]
    )
    rotations = Slerp(
        source_timestamps,
        Rotation.from_matrix(source_poses[:, :3, :3]),
    )(query_timestamps).as_matrix()

    interpolated_poses = np.repeat(
        np.eye(4, dtype=float)[None, :, :], len(query_timestamps), axis=0
    )
    interpolated_poses[:, :3, :3] = rotations
    interpolated_poses[:, :3, 3] = translations
    return WorldTrajectory(
        timestamps=query_timestamps,
        object2world=interpolated_poses,
    )


def create_sensor_trajectory(
    ego_trajectory: WorldTrajectory,
    sensor_to_vehicle: np.ndarray,
) -> WorldTrajectory:
    """Create a sensor-to-world trajectory from vehicle poses and extrinsics."""
    sensor_poses = np.stack(
        [vehicle_pose @ sensor_to_vehicle for vehicle_pose in ego_trajectory.object2world]
    )
    return WorldTrajectory(
        timestamps=ego_trajectory.timestamps,
        object2world=sensor_poses,
    )


def normalize_lidar_intensity(intensities: np.ndarray) -> np.ndarray:
    """
    Normalize LiDAR intensity to [0, 1] range using log-space normalization.
    
    CRITICAL: All LiDAR intensity values must be normalized.
    
    Rules:
    - Non-return (≤0 or non-finite) → -1
    - High outliers (≥p99.9 in raw space) → 1
    - Valid returns → log(raw), then min-max to [0, 1]
    """
    x = intensities.astype(np.float32, copy=True)
    finite_mask = np.isfinite(x)
    valid_raw = x[finite_mask & (x > 0)]
    
    if valid_raw.size == 0:
        x[~finite_mask | (x <= 0)] = -1.0
        return x
    
    # Find 99.9th percentile in raw space
    p99 = float(np.quantile(valid_raw, 0.999))
    fit_vals = valid_raw[valid_raw < p99]
    if fit_vals.size == 0:
        fit_vals = valid_raw
    
    # Log-transform for normalization
    log_fit_vals = np.log(np.clip(fit_vals, 1e-12, None))
    i_min = float(np.min(log_fit_vals))
    i_max = float(np.max(log_fit_vals))
    denom = max(i_max - i_min, 1e-12)
    
    # Apply transformations
    nonreturn = (~finite_mask) | (x <= 0)
    outlier = (~nonreturn) & (x >= p99)
    valid = (~nonreturn) & (~outlier)
    
    if np.any(valid):
        x[valid] = (np.log(np.clip(x[valid], 1e-12, None)) - i_min) / denom
    
    x[nonreturn] = -1.0
    x[outlier] = 1.0
    
    return np.clip(x, -1.0, 1.0)


def create_cameras(
    calib: SensorCalibration,
    dataset: Dict,
    ego_trajectory: WorldTrajectory,
) -> List[Camera]:
    """
    Create Camera objects from dataset.
    
    For each camera:
    1. Load intrinsics and extrinsics
    2. Create ImageData entries for each frame
    3. Link images to disk locations
    """
    cameras = []
    
    # TODO: Get list of cameras in your dataset
    camera_ids = [0]  # Example: one front camera
    
    for cam_id in camera_ids:
        # Parse calibration
        intrinsics = calib.get_camera_intrinsics(cam_id)
        extrinsics = calib.get_camera_extrinsics(cam_id)
        camera_timestamps = np.asarray(
            dataset.get("camera_timestamps", {}).get(cam_id, dataset["timestamps"]),
            dtype=float,
        )
        
        # Create image data entries (one per frame)
        images_data = []
        for frame_idx, timestamp in enumerate(camera_timestamps):
            # TODO: Map frame index to actual image file path
            image_path = f"images/camera_{cam_id}/{frame_idx:04d}.png"
            
            images_data.append(
                ImageData(
                    filepath=image_path,
                    timestamp=timestamp,
                )
            )
        
        # Create Camera object
        camera = Camera(
            id=cam_id,
            name=f"camera_{cam_id}",
            camera_model="opencv",
            distortion_parameters=DistortionParameters(),
            extrinsics=extrinsics,
            intrinsics=intrinsics,
            # Use image-time ego poses, then compose camera-to-vehicle extrinsics.
            trajectory=create_sensor_trajectory(
                interpolate_trajectory(ego_trajectory, camera_timestamps),
                extrinsics,
            ),
            images_data=images_data,
        )
        
        cameras.append(camera)
    
    return cameras


def create_lidars(
    calib: SensorCalibration,
    dataset: Dict,
    ego_trajectory: WorldTrajectory,
) -> List[Lidar]:
    """
    Create LiDAR objects from dataset.
    
    For each LiDAR:
    1. Load intrinsics and extrinsics
    2. Create LidarPointCloudData entries for each scan
    3. Link point clouds to disk locations
    """
    lidars = []
    
    # TODO: Get list of LiDARs in your dataset
    lidar_ids = [0]  # Example: one top-mounted LiDAR
    
    for lidar_id in lidar_ids:
        # Parse calibration
        intrinsics = calib.get_lidar_intrinsics(lidar_id)
        extrinsics = calib.get_lidar_extrinsics(lidar_id)
        
        # Create point cloud data entries (one per scan)
        lidar_pcds_data = []
        for frame_idx, timestamp in enumerate(dataset["timestamps"]):
            # TODO: Map frame index to actual point cloud file path
            pcd_path = f"lidar/lidar_{lidar_id}/{frame_idx:04d}.ply"
            
            lidar_pcds_data.append(
                LidarPointCloudData(
                    filepath=pcd_path,
                    timestamp=timestamp,
                    egomotion_compensated=True,  # Was this unrolled for motion?
                    has_missing_points=False,  # Does this scan have invalid points?
                )
            )
        
        # Create LiDAR object
        lidar = Lidar(
            id=lidar_id,
            name=f"lidar_{lidar_id}",
            extrinsics=extrinsics,
            intrinsics=intrinsics,
            # The trajectory stores LiDAR-to-world poses, not vehicle-to-world.
            trajectory=create_sensor_trajectory(ego_trajectory, extrinsics),
            lidar_pcds_data=lidar_pcds_data,
            use_for_training=True,
        )
        
        lidars.append(lidar)
    
    return lidars


def create_actors(
    dataset: Dict,
    ego_trajectory: WorldTrajectory,
) -> List[Actor]:
    """
    Create Actor objects from annotations.
    
    For each actor in the dataset:
    1. Extract trajectory (positions + orientations over time)
    2. Extract dimensions (length, width, height)
    3. Extract type (vehicle, pedestrian, cyclist, other)
    """
    actors = []
    
    # TODO: Get list of actors/annotations
    if "actors" not in dataset:
        return actors
    
    for actor_data in dataset["actors"]:
        # TODO: Parse actor ID and type
        actor_id = str(actor_data.get("id", 0))
        actor_type = actor_data.get("type", "other")  # Must be one of: vehicle, pedestrian, cyclist, other
        
        # TODO: Parse actor trajectory
        # Should have timestamps matching ego trajectory
        actor_timestamps = np.array(dataset["timestamps"])
        actor_poses = np.zeros((len(actor_timestamps), 4, 4))
        
        # TODO: Populate actor poses for each frame
        
        # TODO: Parse actor dimensions
        actor = Actor(
            id=actor_id,
            type=actor_type,
            dimensions=ActorDimensions(
                length=4.5,   # Forward direction (meters)
                width=2.0,    # Left direction (meters)
                height=1.5,   # Up direction (meters)
            ),
            trajectory=WorldTrajectory(
                timestamps=actor_timestamps,
                object2world=actor_poses,
            ),
        )
        
        actors.append(actor)
    
    return actors


# ============================================================================
# PART 5: Main Conversion Function
# ============================================================================

def convert_dataset(input_path: Path, output_path: Path) -> None:
    """
    Main conversion function.
    
    Orchestrates the conversion from raw dataset format to MIMIC format.
    """
    print(f"Converting dataset from {input_path} to {output_path}")
    
    # Step 1: Load raw dataset
    dataset = load_raw_dataset(input_path)
    print(f"Loaded dataset with {len(dataset['timestamps'])} frames")
    
    # Step 2: Parse calibration
    calib = SensorCalibration(dataset.get("calibration", {}))
    print("Parsed sensor calibration")
    
    # Step 3: Create ego trajectory
    # TODO: Load ego poses from dataset
    ego_timestamps = np.array(dataset["timestamps"])
    ego_poses = np.tile(np.eye(4), (len(ego_timestamps), 1, 1))
    ego_trajectory = create_ego_trajectory(ego_timestamps, ego_poses)
    print(f"Created ego trajectory with {len(ego_trajectory.timestamps)} frames")
    
    # Step 4: Create sensors
    cameras = create_cameras(calib, dataset, ego_trajectory)
    lidars = create_lidars(calib, dataset, ego_trajectory)
    print(f"Created {len(cameras)} cameras, {len(lidars)} LiDARs")
    
    # Step 5: Create actors
    actors = create_actors(dataset, ego_trajectory)
    print(f"Created {len(actors)} actors")
    
    # Step 6: Assemble scene
    scene = Scene(
        ego_vehicle=EgoVehicle(
            trajectory=ego_trajectory,
            sensors=Sensors(
                cameras=cameras,
                lidars=lidars if lidars else None,
                radars=None,
            ),
        ),
        actors=actors,
    )
    print("Assembled scene")
    
    # Step 7: Save scene
    output_path.mkdir(parents=True, exist_ok=True)
    scene_json_path = output_path / "scene.json"
    
    scene_dict = scene.model_dump(mode="json")
    with open(scene_json_path, "w") as f:
        json.dump(scene_dict, f, indent=2)
    
    print(f"✓ Scene saved to {scene_json_path}")
    print(f"  Total file size: {scene_json_path.stat().st_size / 1e6:.1f} MB")


# ============================================================================
# PART 6: CLI Interface
# ============================================================================

def main():
    parser = argparse.ArgumentParser(
        description="Convert custom dataset to MIMIC format"
    )
    parser.add_argument(
        "input_path",
        type=Path,
        help="Path to input dataset",
    )
    parser.add_argument(
        "--output-path",
        type=Path,
        default=None,
        help="Path to output MIMIC dataset (default: ./mimic_output)",
    )
    
    args = parser.parse_args()
    
    if args.output_path is None:
        args.output_path = Path("./mimic_output")
    
    try:
        convert_dataset(args.input_path, args.output_path)
        print("\n✓ Conversion complete!")
    except Exception as e:
        print(f"\n✗ Conversion failed: {e}")
        raise


if __name__ == "__main__":
    main()