"""
将BEVFusion实时输出转换为CARLA可用格式

用于在CARLA仿真中实时显示BEVFusion的检测和地图结果

使用方法:
    python tools/convert_bevfusion_output_to_carla.py \
        --bevfusion-output results/predictions.pkl \
        --output carla_visualization.json
"""

import pickle
import argparse
import json
import numpy as np


class BEVFusionToCarlaConverter:
    """BEVFusion输出转CARLA可视化格式"""
    
    def __init__(self):
        pass
    
    def convert_detection_results(self, bevfusion_output):
        """
        转换3D检测结果
        
        BEVFusion格式:
        {
            'boxes_3d': (N, 9),  # x,y,z,dx,dy,dz,yaw,vx,vy
            'scores_3d': (N,),
            'labels_3d': (N,)
        }
        
        CARLA格式:
        [
            {
                'type': 'vehicle',
                'location': {'x': ..., 'y': ..., 'z': ...},
                'extent': {'x': ..., 'y': ..., 'z': ...},
                'rotation': {'pitch': 0, 'yaw': ..., 'roll': 0},
                'velocity': {'x': ..., 'y': ...},
                'confidence': 0.95
            },
            ...
        ]
        """
        carla_objects = []
        
        boxes = bevfusion_output['boxes_3d']
        scores = bevfusion_output['scores_3d']
        labels = bevfusion_output['labels_3d']
        
        class_names = ['car', 'truck', 'construction_vehicle', 'bus', 'trailer',
                      'barrier', 'motorcycle', 'bicycle', 'pedestrian', 'traffic_cone']
        
        for i in range(len(boxes)):
            x, y, z, dx, dy, dz, yaw, vx, vy = boxes[i]
            
            carla_obj = {
                'type': class_names[labels[i]],
                'location': {
                    'x': float(x),
                    'y': float(y),
                    'z': float(z)
                },
                'extent': {
                    'x': float(dx / 2),
                    'y': float(dy / 2),
                    'z': float(dz / 2)
                },
                'rotation': {
                    'pitch': 0.0,
                    'yaw': float(np.degrees(yaw)),
                    'roll': 0.0
                },
                'velocity': {
                    'x': float(vx),
                    'y': float(vy)
                },
                'confidence': float(scores[i])
            }
            
            carla_objects.append(carla_obj)
        
        return carla_objects
    
    def convert_segmentation_results(self, bevfusion_output):
        """
        转换BEV分割结果
        
        BEVFusion格式:
        {
            'masks_bev': (H, W, C)  # 每个类别的mask
        }
        
        CARLA格式:
        {
            'drivable_area': [[x,y], ...],  # 多边形顶点
            'lanes': [[[x,y], ...], ...],   # 车道线列表
            ...
        }
        """
        masks = bevfusion_output['masks_bev']
        
        # 从mask提取轮廓
        import cv2
        
        carla_segmentation = {}
        
        class_names = ['drivable_area', 'ped_crossing', 'walkway', 
                      'stop_line', 'carpark_area', 'divider']
        
        for i, class_name in enumerate(class_names):
            mask = masks[:, :, i]
            
            # 提取轮廓
            contours, _ = cv2.findContours(
                mask.astype(np.uint8),
                cv2.RETR_EXTERNAL,
                cv2.CHAIN_APPROX_SIMPLE
            )
            
            # 转为坐标点列表
            polygons = []
            for contour in contours:
                if len(contour) > 10:  # 过滤小轮廓
                    # 从像素坐标转为米坐标
                    points = contour.squeeze()
                    points_meter = self.pixel_to_meter(points, masks.shape)
                    polygons.append(points_meter.tolist())
            
            carla_segmentation[class_name] = polygons
        
        return carla_segmentation
    
    def convert_vector_map_results(self, bevfusion_output):
        """
        转换矢量地图结果(MapTR输出)
        
        BEVFusion/MapTR格式:
        {
            'vectors': [
                {
                    'class': 0,  # divider
                    'points': [[x,y], ...],
                    'score': 0.95
                },
                ...
            ]
        }
        
        CARLA格式:
        {
            'lanes': [
                {
                    'points': [[x,y,z], ...],
                    'type': 'divider',
                    'confidence': 0.95
                },
                ...
            ]
        }
        """
        vectors = bevfusion_output.get('vectors', [])
        
        carla_vectors = []
        
        class_map = {0: 'divider', 1: 'boundary', 2: 'ped_crossing'}
        
        for vec in vectors:
            points_2d = np.array(vec['points'])
            
            # 添加z坐标（假设在地面）
            points_3d = np.column_stack([
                points_2d,
                np.zeros(len(points_2d))
            ])
            
            carla_vec = {
                'points': points_3d.tolist(),
                'type': class_map.get(vec['class'], 'unknown'),
                'confidence': vec.get('score', 1.0)
            }
            
            carla_vectors.append(carla_vec)
        
        return {'lanes': carla_vectors}
    
    def pixel_to_meter(self, points_pixel, shape):
        """像素坐标转米坐标"""
        H, W = shape[:2]
        
        # 假设BEV范围 [-50, 50] x [-50, 50]
        points_meter = points_pixel.copy().astype(float)
        points_meter[:, 0] = (points_pixel[:, 0] / W) * 100 - 50
        points_meter[:, 1] = (points_pixel[:, 1] / H) * 100 - 50
        
        return points_meter
    
    def convert_full_output(self, bevfusion_output, output_path):
        """转换完整的BEVFusion输出"""
        carla_format = {}
        
        # 1. 检测结果
        if 'boxes_3d' in bevfusion_output:
            carla_format['objects'] = self.convert_detection_results(bevfusion_output)
        
        # 2. 分割结果
        if 'masks_bev' in bevfusion_output:
            carla_format['segmentation'] = self.convert_segmentation_results(bevfusion_output)
        
        # 3. 矢量地图
        if 'vectors' in bevfusion_output or 'vector_map' in bevfusion_output:
            vec_output = bevfusion_output.get('vector_map', bevfusion_output)
            carla_format['vector_map'] = self.convert_vector_map_results(vec_output)
        
        # 保存JSON
        with open(output_path, 'w') as f:
            json.dump(carla_format, f, indent=2)
        
        print(f"✅ CARLA格式文件已保存: {output_path}")
        
        return carla_format


def main():
    parser = argparse.ArgumentParser(description='BEVFusion输出转CARLA格式')
    parser.add_argument('--bevfusion-output', type=str, required=True,
                       help='BEVFusion预测结果文件(.pkl)')
    parser.add_argument('--output', type=str, required=True,
                       help='输出CARLA格式文件(.json)')
    
    args = parser.parse_args()
    
    # 加载BEVFusion输出
    print(f"加载BEVFusion输出: {args.bevfusion_output}")
    with open(args.bevfusion_output, 'rb') as f:
        bevfusion_output = pickle.load(f)
    
    # 转换
    converter = BEVFusionToCarlaConverter()
    converter.convert_full_output(bevfusion_output, args.output)
    
    print(f"\n在CARLA中使用:")
    print(f"  import carla")
    print(f"  # 加载JSON并可视化")
    print(f"  with open('{args.output}') as f:")
    print(f"      data = json.load(f)")
    print(f"  # 在CARLA中绘制...")


if __name__ == '__main__':
    main()