# get_flops_smp.py
import torch
import argparse
from fvcore.nn import FlopCountAnalysis, flop_count_table

# 同样，需要能够导入 train_smp.py 中的函数
from train import initialize_components
import config

# TODO 使用方法 TODO
# python Tool_get_params_and_FLOPs.py -a Unet --shape 512 512

def parse_args():
    parser = argparse.ArgumentParser(description='Get FLOPs and Params of an SMP model')
    parser.add_argument(
        '-a', '--architecture',
        type=str,
        required=True,
        choices=list(config.ALL_MODEL_CONFIGS.keys()),
        help="The model architecture to analyze."
    )
    parser.add_argument(
        '--shape',
        type=int,
        nargs='+',
        default=[512, 512],
        help='Input image size for analysis, e.g., --shape 512 512'
    )
    args = parser.parse_args()
    return args

def main():
    args = parse_args()
    
    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.')
        
    # 1. 初始化模型
    print(f"Initializing model: '{args.architecture}'...")
    num_classes = len(config.CLASSES)
    model, _, _, _ = initialize_components(args.architecture, num_classes, config.SEG_MODE)

    # 如果模型在训练时使用了 DataParallel，需要移除
    if isinstance(model, torch.nn.DataParallel):
         model = model.module
         
    model.to(config.DEVICE)
    model.eval()

    # 2. 创建伪造输入
    dummy_input = torch.randn(1, 3, h, w).to(config.DEVICE)
    
    # 3. 计算参数量
    # 分别计算总参数和可训练参数
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    
    # 4. 使用 fvcore 计算 FLOPs
    print("Analyzing model FLOPs... (This may take a moment)")
    flops = FlopCountAnalysis(model, dummy_input)

    # 5. 打印结果
    split_line = '=' * 30
    print(f"\n{split_line}")
    print(f"Model: {args.architecture}")
    print(f"Input shape: (1, 3, {h}, {w})")
    print(f"{split_line}")
    
    # 使用 fvcore 的表格打印工具，非常清晰
    print(flop_count_table(flops, max_depth=4))
    
    # 打印总结
    gflops = flops.total() / 1e9
    total_params_m = total_params / 1e6
    trainable_params_m = trainable_params / 1e6
    
    print(f"{split_line}")
    print(f"Total FLOPs: {gflops:.2f} GFLOPs")
    print(f"Total Parameters: {total_params_m:.2f} M")
    print(f"Trainable Parameters: {trainable_params_m:.2f} M")
    print(f"{split_line}")
    print('!!! Please be cautious: FLOPs computation may not be perfectly accurate for all custom ops.')

if __name__ == '__main__':
    main()