import argparse
import time
import torch
from mmengine import Config
from mmengine.config import DictAction
from mmengine.device import get_device
from mmengine.registry import init_default_scope
from mmengine.runner import Runner, autocast, load_checkpoint
from mmdet3d.registry import MODELS
from tools.misc.fuse_conv_bn import fuse_module

def parse_args():
    parser = argparse.ArgumentParser(description='MMDet benchmark a model')
    parser.add_argument('config', help='test config file path')
    parser.add_argument('checkpoint', help='checkpoint file')
    parser.add_argument('--samples', default=2000, type=int, help='samples to benchmark')
    parser.add_argument('--log-interval', default=10, type=int, help='interval of logging')
    parser.add_argument('--amp', action='store_true', help='Use AMP')
    parser.add_argument('--fuse-conv-bn', action='store_true', help='Fuse conv and bn')
    # 补丁 1：增加对 cfg-options 的支持
    parser.add_argument('--cfg-options', nargs='+', action=DictAction, help='override config')
    args = parser.parse_args()
    return args

def main():
    args = parse_args()
    init_default_scope('mmdet3d')
    cfg = Config.fromfile(args.config)
    
    # 补丁 2：将命令行参数合并到配置中
    if args.cfg_options is not None:
        cfg.merge_from_dict(args.cfg_options)

    if cfg.env_cfg.get('cudnn_benchmark', False):
        torch.backends.cudnn.benchmark = True

    dataloader = Runner.build_dataloader(cfg.test_dataloader)
    model = MODELS.build(cfg.model)
    load_checkpoint(model, args.checkpoint, map_location='cpu')
    if args.fuse_conv_bn:
        model = fuse_module(model)
    model.to(get_device())
    model.eval()

    num_warmup = 5
    pure_inf_time = 0
    total_batches = len(dataloader)

    for i, data in enumerate(dataloader):
        # ================= [新增的打印 Shape 代码] =================
        if i == 0:
            print("\n" + "🔥" * 25)
            print("🔍 [Debug] 当前输入模型的 Data Shape:")
            if 'inputs' in data:
                inputs = data['inputs']
                # 1. 打印雷达点云的 shape
                if 'points' in inputs:
                    print(f"   👉 雷达点云 (points): 包含 {len(inputs['points'])} 个样本 (Batch Size)")
                    for b_idx, pts in enumerate(inputs['points']):
                        print(f"      - 样本 {b_idx} shape: {pts.shape}")
                
                # 2. 打印相机图像的 shape (如果跑的是多模态)
                if 'img' in inputs:
                    print(f"   👉 相机图像 (img): {inputs['img'][0].shape}")
            print("🔥" * 25 + "\n")
        # =========================================================
        torch.cuda.synchronize()
        start_time = time.perf_counter()

        with autocast(enabled=args.amp):
            model.test_step(data)

        torch.cuda.synchronize()
        elapsed = time.perf_counter() - start_time

        if i >= num_warmup:
            pure_inf_time += elapsed
            if (i + 1) % args.log_interval == 0:
                fps = (i + 1 - num_warmup) / pure_inf_time
                print(f'Done batch [{i + 1:<3}/ {total_batches}], speed: {fps:.1f} batch / s')

        # 补丁 3：只要读完当前数据集，立刻退出并打印，防止哑巴
        if (i + 1) == args.samples or (i + 1) == total_batches:
            break

    if (i + 1) > num_warmup:
        fps = (i + 1 - num_warmup) / pure_inf_time
        bs = cfg.test_dataloader.get('batch_size', 1)
        print(f'\n✅ 测试完成！')
        print(f'🚀 纯模型前向速度 (Batch/s): {fps:.1f}')
        print(f'🚀 真实吞吐量 (等效 task/s): {fps * bs:.1f} frames / s')
    else:
        print('数据量太少，连 Warmup 都没跑完。')

if __name__ == '__main__':
    main()
