#!/usr/bin/env python
"""
单GPU测试脚本，用于验证修复后的配置
"""
import os
import sys
import torch
import mmcv
from mmcv import Config
from mmdet3d.apis import single_gpu_test
from mmdet3d.datasets import build_dataloader, build_dataset
from mmdet3d.models import build_model
from mmcv.runner import load_checkpoint

# 设置环境
os.environ['PATH'] = '/opt/conda/bin:' + os.environ.get('PATH', '')
os.environ['LD_LIBRARY_PATH'] = '/opt/conda/lib/python3.8/site-packages/torch/lib:/opt/conda/lib:/usr/local/cuda/lib64:' + os.environ.get('LD_LIBRARY_PATH', '')
os.environ['PYTHONPATH'] = '/workspace/bevfusion:' + os.environ.get('PYTHONPATH', '')

def main():
    # 加载配置
    config_file = 'configs/nuscenes/det/transfusion/secfpn/camera+lidar/swint_v0p075/multitask_BEV2X_phase4b_rmtppad_segmentation.yaml'
    cfg = Config.fromfile(config_file)

    # 设置测试模式
    cfg.data.val.test_mode = True

    # 构建数据集
    print("构建数据集...")
    from mmdet3d.datasets import NuScenesDataset

    val_config = {
        'type': 'NuScenesDataset',
        'dataset_root': cfg.dataset_root,
        'ann_file': cfg.dataset_root + 'nuscenes_infos_val.pkl',
        'pipeline': cfg.evaluation.pipeline,
        'object_classes': cfg.object_classes,
        'map_classes': cfg.map_classes,
        'modality': {
            'use_lidar': True,
            'use_camera': True,
            'use_radar': False,
            'use_map': False,
            'use_external': False
        },
        'test_mode': True,
        'use_valid_flag': False,
        'box_type_3d': 'LiDAR',
        'load_interval': cfg.data.val.load_interval if hasattr(cfg.data.val, 'load_interval') else 1
    }

    dataset = build_dataset(val_config)
    print(f"数据集大小: {len(dataset)}")

    # 构建数据加载器
    data_loader = build_dataloader(
        dataset,
        samples_per_gpu=1,
        workers_per_gpu=0,
        dist=False,
        shuffle=False,
    )

    # 构建模型
    print("构建模型...")
    cfg.model.train_cfg = None
    model = build_model(cfg.model, test_cfg=cfg.get("test_cfg"))
    print(f"模型类型: {type(model)}")
    print(f"模型heads: {list(model.heads.keys())}")

    # 加载checkpoint
    checkpoint_path = 'runs/run-4c8ec7e5-fabdc997/epoch_1.pth'
    print(f"加载checkpoint: {checkpoint_path}")
    checkpoint = load_checkpoint(model, checkpoint_path, map_location="cpu")

    if "CLASSES" in checkpoint.get("meta", {}):
        model.CLASSES = checkpoint["meta"]["CLASSES"]
    else:
        model.CLASSES = dataset.CLASSES

    # 单GPU测试
    print("开始单GPU测试...")
    model = model.cuda()
    outputs = single_gpu_test(model, data_loader)

    print(f"测试完成，结果数量: {len(outputs)}")

    # 检查结果
    if outputs:
        sample = outputs[0]
        print("第一个样本的结果结构:")
        for key, value in sample.items():
            if hasattr(value, 'shape'):
                print(f"  {key}: {type(value)}, shape: {value.shape}")
            elif hasattr(value, '__len__') and not isinstance(value, str):
                print(f"  {key}: {type(value)}, length: {len(value)}")
            else:
                print(f"  {key}: {type(value)}, value: {value}")

        # 检查检测结果
        if 'boxes_3d' in sample:
            print(f"检测框数量: {len(sample['boxes_3d'])}")
        if 'scores_3d' in sample:
            print(f"分数形状: {sample['scores_3d'].shape}")

if __name__ == '__main__':
    main()