"""
将nuScenes/MapTR矢量地图转换为CARLA OpenDRIVE格式

使用方法:
    python tools/convert_to_carla_opendrive.py \
        --vector-maps data/nuscenes/vector_maps_bevfusion.pkl \
        --nuscenes-root data/nuscenes \
        --output carla_maps/nuscenes_map.xodr
"""

import pickle
import argparse
import numpy as np
from lxml import etree
from pathlib import Path
from tqdm import tqdm
from collections import defaultdict
from pyquaternion import Quaternion


class NuScenesToOpenDRIVE:
    """nuScenes矢量地图转CARLA OpenDRIVE格式"""
    
    def __init__(self, vector_maps_path, nuscenes_root, bev_range=50):
        # 加载矢量地图
        print(f"加载矢量地图: {vector_maps_path}")
        with open(vector_maps_path, 'rb') as f:
            self.vector_maps = pickle.load(f)
        
        print(f"加载样本数: {len(self.vector_maps)}")
        
        # 加载nuScenes数据以获取全局位姿
        from nuscenes.nuscenes import NuScenes
        self.nusc = NuScenes(version='v1.0-trainval', dataroot=nuscenes_root, verbose=False)
        
        self.bev_range = bev_range
        
        # 加载训练info以获取token映射
        info_path = f"{nuscenes_root}/nuscenes_infos_train.pkl"
        with open(info_path, 'rb') as f:
            infos = pickle.load(f)
        self.infos = infos['infos'] if isinstance(infos, dict) else infos
    
    def convert(self, output_path):
        """
        主转换函数
        
        流程:
        1. 将所有局部矢量转为全局坐标
        2. 聚合重叠的矢量（去重）
        3. 构建OpenDRIVE XML
        """
        print("\n开始转换为OpenDRIVE格式...")
        
        # 1. 转为全局坐标
        print("步骤1: 转换为全局坐标...")
        global_vectors = self.convert_to_global_coords()
        
        # 2. 聚合矢量
        print("步骤2: 聚合和去重...")
        merged_vectors = self.merge_vectors(global_vectors)
        
        # 3. 构建road网络
        print("步骤3: 构建道路网络...")
        roads = self.build_road_network(merged_vectors)
        
        # 4. 生成OpenDRIVE XML
        print("步骤4: 生成OpenDRIVE XML...")
        self.generate_opendrive_xml(roads, output_path)
        
        print(f"\n✅ OpenDRIVE文件已保存: {output_path}")
        
        # 统计
        total_roads = len(roads)
        total_lanes = sum(len(r['lanes']) for r in roads)
        print(f"\n统计:")
        print(f"  Roads: {total_roads}")
        print(f"  Lanes: {total_lanes}")
    
    def convert_to_global_coords(self):
        """将所有局部矢量转换为全局坐标"""
        global_vectors = {
            'divider': [],
            'boundary': [],
            'ped_crossing': []
        }
        
        for info in tqdm(self.infos, desc="转换坐标"):
            token = info['token']
            
            if token not in self.vector_maps:
                continue
            
            vectors = self.vector_maps[token]
            
            # 获取全局位姿
            try:
                sample = self.nusc.get('sample', token)
                lidar_token = sample['data']['LIDAR_TOP']
                lidar_data = self.nusc.get('sample_data', lidar_token)
                ego_pose = self.nusc.get('ego_pose', lidar_data['ego_pose_token'])
            except:
                continue
            
            # 转换每个类别的矢量
            for category in ['divider', 'boundary', 'ped_crossing']:
                for vec in vectors.get(category, []):
                    points_norm = np.array(vec['points'])  # (20, 2) 归一化坐标
                    
                    # 反归一化到ego坐标
                    points_ego = self.denormalize(points_norm)
                    
                    # 转全局坐标
                    points_global = self.ego_to_global(points_ego, ego_pose)
                    
                    global_vectors[category].append({
                        'points': points_global.tolist(),
                        'token': token
                    })
        
        return global_vectors
    
    def denormalize(self, points_norm):
        """反归一化: [0,1] → [-50,50]米"""
        points = points_norm * (2 * self.bev_range) - self.bev_range
        return points
    
    def ego_to_global(self, points_ego, ego_pose):
        """ego坐标系 → 全局坐标系"""
        # 旋转
        quat = Quaternion(ego_pose['rotation'])
        rot_matrix = quat.rotation_matrix[:2, :2]
        points_rotated = points_ego @ rot_matrix.T
        
        # 平移
        translation = np.array(ego_pose['translation'][:2])
        points_global = points_rotated + translation
        
        return points_global
    
    def merge_vectors(self, global_vectors, distance_threshold=2.0):
        """聚合相近的矢量（去重）"""
        merged = {
            'divider': [],
            'boundary': [],
            'ped_crossing': []
        }
        
        for category in ['divider', 'boundary', 'ped_crossing']:
            vectors = global_vectors[category]
            
            if len(vectors) == 0:
                continue
            
            print(f"  聚合{category}: {len(vectors)}个矢量...")
            
            # 简单去重：距离很近的矢量只保留一个
            used = [False] * len(vectors)
            
            for i in range(len(vectors)):
                if used[i]:
                    continue
                
                current = np.array(vectors[i]['points'])
                used[i] = True
                
                # 找到相似的矢量并平均
                similar_indices = [i]
                
                for j in range(i+1, len(vectors)):
                    if used[j]:
                        continue
                    
                    other = np.array(vectors[j]['points'])
                    
                    # 计算Hausdorff距离
                    dist = self.hausdorff_distance(current, other)
                    
                    if dist < distance_threshold:
                        similar_indices.append(j)
                        used[j] = True
                
                # 平均相似的矢量
                if len(similar_indices) > 1:
                    all_points = [np.array(vectors[idx]['points']) for idx in similar_indices]
                    merged_points = np.mean(all_points, axis=0)
                else:
                    merged_points = current
                
                merged[category].append({
                    'points': merged_points.tolist()
                })
        
        print(f"  去重后: divider={len(merged['divider'])}, "
              f"boundary={len(merged['boundary'])}, "
              f"crossing={len(merged['ped_crossing'])}")
        
        return merged
    
    def hausdorff_distance(self, points1, points2):
        """计算Hausdorff距离"""
        from scipy.spatial.distance import cdist
        
        dist_matrix = cdist(points1, points2)
        d1 = dist_matrix.min(axis=1).max()
        d2 = dist_matrix.min(axis=0).max()
        
        return max(d1, d2)
    
    def build_road_network(self, merged_vectors):
        """构建道路网络"""
        roads = []
        
        # 简化版：每个divider对应一条road
        # 实际应该做拓扑分析，连接成road network
        
        for idx, divider in enumerate(merged_vectors['divider']):
            points = np.array(divider['points'])
            
            if len(points) < 2:
                continue
            
            # 计算road的长度
            diffs = np.diff(points, axis=0)
            lengths = np.sqrt((diffs**2).sum(axis=1))
            total_length = lengths.sum()
            
            # 计算初始heading
            dx = points[1, 0] - points[0, 0]
            dy = points[1, 1] - points[0, 1]
            hdg = np.arctan2(dy, dx)
            
            roads.append({
                'id': idx,
                'length': total_length,
                'start': points[0],
                'hdg': hdg,
                'geometry': points,
                'lanes': [
                    {
                        'id': 1,
                        'type': 'driving',
                        'level': 0,
                        'width': 3.5  # 标准车道宽度
                    }
                ]
            })
        
        return roads
    
    def generate_opendrive_xml(self, roads, output_path):
        """生成OpenDRIVE XML文件"""
        # 创建根元素
        root = etree.Element('OpenDRIVE')
        
        # Header
        header = etree.SubElement(root, 'header',
            revMajor='1',
            revMinor='6',
            name='nuScenes Map Converted',
            version='1.00',
            date='2025-10-18',
            north='0.0',
            south='0.0',
            east='0.0',
            west='0.0'
        )
        
        # Roads
        for road_data in roads:
            road = self.create_road_xml(road_data)
            root.append(road)
        
        # 格式化并保存
        tree = etree.ElementTree(root)
        
        # 确保目录存在
        Path(output_path).parent.mkdir(parents=True, exist_ok=True)
        
        tree.write(
            output_path,
            pretty_print=True,
            xml_declaration=True,
            encoding='UTF-8'
        )
    
    def create_road_xml(self, road_data):
        """创建单个road的XML元素"""
        road = etree.Element('road',
            name=f"road_{road_data['id']}",
            length=f"{road_data['length']:.2f}",
            id=str(road_data['id']),
            junction='-1'
        )
        
        # 1. Plan View (几何)
        plan_view = etree.SubElement(road, 'planView')
        
        # 简化：使用line几何（实际应该拟合曲线）
        points = road_data['geometry']
        
        # 分段geometry
        s = 0.0
        for i in range(len(points) - 1):
            x, y = points[i]
            dx = points[i+1, 0] - x
            dy = points[i+1, 1] - y
            length = np.sqrt(dx**2 + dy**2)
            hdg = np.arctan2(dy, dx)
            
            geometry = etree.SubElement(plan_view, 'geometry',
                s=f"{s:.2f}",
                x=f"{x:.2f}",
                y=f"{y:.2f}",
                hdg=f"{hdg:.4f}",
                length=f"{length:.2f}"
            )
            
            # Line类型
            line = etree.SubElement(geometry, 'line')
            
            s += length
        
        # 2. Lanes
        lanes = etree.SubElement(road, 'lanes')
        
        # Lane Section
        lane_section = etree.SubElement(lanes, 'laneSection', s='0.0')
        
        # Center lane
        center = etree.SubElement(lane_section, 'center')
        center_lane = etree.SubElement(center, 'lane', id='0', type='none', level='0')
        
        # Right lanes
        right = etree.SubElement(lane_section, 'right')
        for lane_info in road_data['lanes']:
            lane = etree.SubElement(right, 'lane',
                id=str(lane_info['id']),
                type=lane_info['type'],
                level=str(lane_info['level'])
            )
            
            # Lane width
            width = etree.SubElement(lane, 'width',
                sOffset='0.0',
                a=str(lane_info['width']),
                b='0.0',
                c='0.0',
                d='0.0'
            )
        
        return road


def main():
    parser = argparse.ArgumentParser(description='转换nuScenes矢量地图为OpenDRIVE')
    parser.add_argument('--vector-maps', type=str, required=True,
                       help='矢量地图pkl文件路径')
    parser.add_argument('--nuscenes-root', type=str, default='data/nuscenes',
                       help='nuScenes数据根目录')
    parser.add_argument('--output', type=str, required=True,
                       help='输出OpenDRIVE文件路径(.xodr)')
    parser.add_argument('--bev-range', type=float, default=50,
                       help='BEV范围（米）')
    
    args = parser.parse_args()
    
    # 转换
    converter = NuScenesToOpenDRIVE(
        args.vector_maps,
        args.nuscenes_root,
        args.bev_range
    )
    
    converter.convert(args.output)
    
    print(f"\n✅ 转换完成！")
    print(f"OpenDRIVE文件: {args.output}")
    print(f"\n在CARLA中使用:")
    print(f"  1. 复制到: CARLA/Import/{Path(args.output).name}")
    print(f"  2. 启动CARLA并导入地图")


if __name__ == '__main__':
    main()