# benchmark_smp.py
import torch
import time
import argparse
import numpy as np
from pathlib import Path
from mmengine.model.utils import revert_sync_batchnorm

# 假设 train.py 在同一目录下，以便导入其函数
# 如果不在，请确保路径正确
from train import initialize_components
import config

# TODO 使用方法 TODO
# 交互模式（推荐）
# python Tool_benchmark_smp.py -a Unet --shape 512 512
# 非交互模式（用于脚本）
# python Tool_benchmark_smp.py -a Unet -c "outputs/Unet_2025-10-13_19-30-00/best_model.pth" --shape 512 512

# --- Function copied from predict.py ---
def select_trained_model_path(base_dir: Path, model_architecture: str) -> Path:
    """
    查找指定架构的所有训练运行目录，并让用户选择一个。
    """
    print(f"\nSearching for '{model_architecture}' model runs in '{base_dir}'...")
    
    # 查找所有以架构名开头的文件夹
    run_dirs = sorted([d for d in base_dir.iterdir() if d.is_dir() and d.name.startswith(model_architecture+'_')])
    
    if not run_dirs:
        print(f"ERROR: No trained models found for architecture '{model_architecture}'. Please run train_smp.py first.")
        return None

    print("\nPlease select a trained model to benchmark:")
    for i, dir_path in enumerate(run_dirs):
        print(f"  [{i+1}] {dir_path.name}")
    
    while True:
        try:
            choice = input(f"Enter selection (1-{len(run_dirs)}) or press Enter to cancel: ")
            if not choice:
                return None
            choice_idx = int(choice) - 1
            if 0 <= choice_idx < len(run_dirs):
                selected_dir = run_dirs[choice_idx]
                print(f"Selected model: {selected_dir.name}")
                return selected_dir
            else:
                print("Invalid selection. Please try again.")
        except (ValueError, IndexError):
            print("Invalid input. Please enter a number.")

def parse_args():
    parser = argparse.ArgumentParser(description='Benchmark a segmentation model from SMP')
    parser.add_argument(
        '-a', '--architecture',
        type=str,
        required=True,
        choices=list(config.ALL_MODEL_CONFIGS.keys()),
        help="The model architecture to benchmark."
    )
    parser.add_argument(
        '-c', '--checkpoint',
        type=str,
        required=False, # <-- Changed to False
        help='(Optional) Path to the checkpoint file. If omitted, an interactive selection will be shown.'
    )
    parser.add_argument(
        '--shape',
        type=int,
        nargs='+',
        default=[512, 512],
        help='Input image size for benchmarking, e.g., --shape 512 512'
    )
    parser.add_argument(
        '--log-interval',
        type=int,
        default=50,
        help='Interval of logging.'
    )
    parser.add_argument(
        '--repeat-times',
        type=int,
        default=3,
        help='Number of times to repeat the benchmark for averaging.'
    )
    args = parser.parse_args()
    return args

def main():
    args = parse_args()

    # --- New logic to select checkpoint ---
    checkpoint_path = None
    if args.checkpoint:
        checkpoint_path = Path(args.checkpoint)
    else:
        # Assumes your outputs are saved in the directory specified in config.py
        # This is consistent with train_smp.py and predict.py
        selected_run_dir = select_trained_model_path(Path(config.PREDICT_BEST_MODEL_DIR), args.architecture)
        if not selected_run_dir:
            print("No model selected. Exiting.")
            sys.exit()
        
        checkpoint_path = selected_run_dir / config.BEST_MODEL_SAVE_NAME
    
    if not checkpoint_path.exists():
        print(f"ERROR: Checkpoint file not found at '{checkpoint_path}'")
        sys.exit()
    # --- End of new logic ---
    
    if len(args.shape) == 1:
        h, w = args.shape[0], args.shape[0]
    elif len(args.shape) == 2:
        h, w = args.shape[0], args.shape[1]
    else:
        raise ValueError('Invalid input shape. Use one or two integers.')

    print(f"\n--- Model Benchmarking ---")
    print(f"Architecture: {args.architecture}")
    print(f"Checkpoint:   {checkpoint_path}")
    print(f"Input Shape:  (1, 3, {h}, {w})")
    print(f"Device:       {config.DEVICE}")
    print("-" * 28)

    # 1. 初始化模型
    num_classes = len(config.CLASSES)
    model, _, _, _ = initialize_components(args.architecture, num_classes, config.SEG_MODE)

    # 2. 加载权重
    print("Loading checkpoint...")
    state_dict = torch.load(checkpoint_path, map_location=config.DEVICE)
    
    # Handle DataParallel saved models by removing 'module.' prefix
    is_dataparallel = any(key.startswith('module.') for key in state_dict.keys())
    if is_dataparallel:
        print("DataParallel model detected, removing 'module.' prefix.")
        new_state_dict = {key.replace('module.', ''): value for key, value in state_dict.items()}
        model.load_state_dict(new_state_dict)
    else:
        model.load_state_dict(state_dict)
    
    model.to(config.DEVICE)
    model.eval()

    # 3. 准备伪造输入数据
    dummy_input = torch.randn(1, 3, h, w).to(config.DEVICE)

    # 4. 开始基准测试
    overall_fps_list = []
    for run_index in range(args.repeat_times):
        print(f"\n>>> Running benchmark iteration {run_index + 1}/{args.repeat_times}...")
        
        num_warmup = 5
        pure_inf_time = 0
        total_iters = 200

        with torch.no_grad():
            for i in range(total_iters + num_warmup):
                if torch.cuda.is_available():
                    torch.cuda.synchronize()
                start_time = time.perf_counter()

                _ = model(dummy_input)

                if torch.cuda.is_available():
                    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 image [{i + 1 - num_warmup:<4}/ {total_iters}], '
                              f'current FPS: {fps:.2f} img/s')
        
        overall_fps = total_iters / pure_inf_time
        overall_fps_list.append(overall_fps)
        print(f'Overall FPS for this run: {overall_fps:.2f} img/s')

    # 5. 总结结果
    print("\n--- Benchmark Summary ---")
    avg_fps = np.mean(overall_fps_list)
    var_fps = np.var(overall_fps_list)
    print(f'Average FPS of {args.repeat_times} runs: {avg_fps:.2f} img/s')
    print(f'FPS Variance of {args.repeat_times} runs: {var_fps:.4f}')

if __name__ == '__main__':
    main()