import os
import glob
import logging
import argparse
import csv
from typing import Dict, Optional, Tuple, List
import time
import numpy as np
from pathlib import Path
import sys

import torch
from fvcore.nn import FlopCountAnalysis
import segmentation_models_pytorch as smp
import config  # Assumes config.py is in the same directory or accessible

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# TODO 使用方法 TODO
# 交互模式（推荐）
# python 2_predict_params_and_FLOPs.py
# 非交互模式（用于脚本）
# python 2_predict_params_and_FLOPs_V2.py --input_dir '../BestMode_Predict_Results_DataSet_Public'
# python 2_predict_params_and_FLOPs_V2.py --shape 512 512 # --shape 1080 1920

def load_model(model_architecture: str, checkpoint_path: Optional[Path]) -> Optional[torch.nn.Module]:
    """
    根据 config.py 直接重建模型架构，并从 checkpoint 路径加载权重。
    此函数不依赖 train.py。
    """
    # 1. 根据 config 动态创建基础模型架构
    logging.info(f"Rebuilding model architecture: '{model_architecture}'")
    try:
        model_params = config.ALL_MODEL_CONFIGS[model_architecture].copy()
        model_class = getattr(smp, model_architecture)
    except (AttributeError, KeyError):
        logging.error(f"Model '{model_architecture}' not found in smp library or config.py!")
        return None

    # 2. 准备初始化参数
    num_classes = len(config.CLASSES)
    params = model_params
    params['classes'] = num_classes
    # 加载权重时，不应再次下载预训练权重
    params['encoder_weights'] = None

    # 自动检测 ViT 编码器以设置特殊参数
    encoder_name = model_params.get('encoder_name', '')
    if 'vit' in encoder_name.lower():
        params['dynamic_img_size'] = True
        logging.info(f"ViT encoder ('{encoder_name}') detected. Setting dynamic_img_size=True.")

    # 3. 创建模型实例
    model = model_class(**params).to(config.DEVICE)

    # 4. 如果提供了有效的 checkpoint 路径，则加载权重
    if checkpoint_path and checkpoint_path.exists():
        logging.info(f"Loading checkpoint: {checkpoint_path}")
        try:
            state_dict = torch.load(checkpoint_path, map_location=torch.device(config.DEVICE))
            # 处理 DataParallel 保存的权重
            is_dataparallel = any(key.startswith('module.') for key in state_dict.keys())
            if is_dataparallel:
                logging.info("DataParallel weights 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)
            logging.info("Checkpoint loaded successfully.")
        except Exception as e:
            logging.error(f"Failed to load checkpoint. Using model with random weights. Error: {e}")
    else:
        logging.info("No valid checkpoint provided. Using model with random weights.")
    
    model.eval()
    return model

# --- 核心功能函数 ---
def find_latest_model_run_path(base_dir: Path, model_architecture: str) -> Optional[Path]:
    """
    非交互式地查找并返回给定架构的最新模型运行目录。
    """
    logging.info(f"Searching for latest '{model_architecture}' model run 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:
        logging.info(f"No trained model folder found for '{model_architecture}'.")
        return None

    # 排序后的最后一个即为最新的
    latest_run_dir = run_dirs[-1]
    logging.info(f"Automatically selected latest model run: {latest_run_dir.name}")
    return latest_run_dir

def calculate_flops_params(model: torch.nn.Module, shape: Tuple[int, int]) -> Optional[Dict[str, str]]:
    """
    为已初始化的模型计算FLOPs和参数量。
    """
    logging.info(f"Analyzing model Params/FLOPs with input shape {shape}...")
    model.to(config.DEVICE)
    model.eval()
    dummy_input = torch.randn(1, 3, shape[0], shape[1]).to(config.DEVICE)
    try:
        flops_analyzer = FlopCountAnalysis(model, dummy_input)
        gflops = flops_analyzer.total() / 1e9
        m_params = sum(p.numel() for p in model.parameters()) / 1e6
        logging.info(f"✅ Params/FLOPs Analysis successful: {gflops:.2f} GFLOPs, {m_params:.2f} M Params")
        return {'flops': f"{gflops:.2f} G", 'params': f"{m_params:.2f} M"}
    except Exception as e:
        logging.error(f"❌ Failed to analyze FLOPs/Params. Error: {e}")
        return None

def benchmark_model(model: torch.nn.Module, shape: Tuple[int, int]) -> Optional[Dict[str, float]]:
    """
    对加载了权重的模型进行FPS基准测试。
    """
    logging.info("Starting FPS benchmark...")
    model.to(config.DEVICE)
    model.eval()
    dummy_input = torch.randn(1, 3, shape[0], shape[1]).to(config.DEVICE)

    overall_fps_list = []
    repeat_times = 3
    
    with torch.no_grad():
        for run_index in range(repeat_times):
            num_warmup = 5
            pure_inf_time = 0
            total_iters = 100

            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 pure_inf_time == 0:
                pure_inf_time = 1e-6
            
            overall_fps = total_iters / pure_inf_time
            overall_fps_list.append(overall_fps)
            print(f"  Run {run_index + 1}/{repeat_times}, FPS: {overall_fps:.2f} img/s")

    if not overall_fps_list:
        return None
        
    avg_fps = np.mean(overall_fps_list)
    fps_variance = np.var(overall_fps_list)
    logging.info(f"✅ FPS Benchmark successful: Average FPS: {avg_fps:.2f} img/s")
    return {'avg_fps': avg_fps, 'fps_variance': fps_variance}

def get_shape_from_path(path: str) -> Optional[Tuple[int, int]]:
    """
    从文件夹路径中提取分辨率。
    """
    # V1. 采用图片本身的尺寸作为输入
    import re
    match = re.search(r'(\d+)x(\d+)', os.path.basename(path))
    if match:
        width, height = int(match.group(1)), int(match.group(2))
        return (height, width)  # Return as H, W
    return None

    # V2. 采用文件夹config中的尺寸作为输入
    # width, height =  config.IMAGE_WIDTH, config.IMAGE_HEIGHT

    if width > 0 and height > 0:
        return (height, width)
    else:
        return None

# --- 主函数 ---
def main(args):
    """
    脚本主入口。
    """
    input_root = args.input_dir
    if not os.path.isdir(input_root):
        logging.error(f"Input directory does not exist: {input_root}")
        return

    # 1. 交互式选择数据集和算法
    existing_dataset_dirs = sorted([d for d in glob.glob(os.path.join(input_root, '*_outputs-SegModel')) if os.path.isdir(d)])
    if not existing_dataset_dirs:
        logging.error(f"No valid dataset folders found in {input_root}.")
        return

    dataset_map = {str(i + 1): path for i, path in enumerate(existing_dataset_dirs)}
    print("\n" + "="*50)
    print("--- Step 1: Please select the dataset to process ---")
    for key, path in dataset_map.items():
        print(f"{key}: {os.path.basename(path)}")
    print("="*50)
    choice1 = input("Enter the dataset number and press Enter: ").strip()
    
    model_archs_to_process = []
    selected_dataset_dir = None
    if choice1 in dataset_map:
        selected_dataset_dir = dataset_map[choice1]
        logging.info(f"You have selected the dataset: [{os.path.basename(selected_dataset_dir)}]")

        alg_dirs_full_path = sorted([d for d in glob.glob(os.path.join(selected_dataset_dir, '*')) if os.path.isdir(d)])
        valid_alg_archs = []
        for path in alg_dirs_full_path:
            base_name = os.path.basename(path).split('_')[0]
            if base_name in config.ALL_MODEL_CONFIGS and base_name not in valid_alg_archs:
                valid_alg_archs.append(base_name)
        
        if valid_alg_archs:
            alg_map = {str(i + 1): name for i, name in enumerate(valid_alg_archs)}
            print("\n" + "="*50)
            print("--- Step 2: Please select the algorithm to process ---")
            print("0: Process [ALL] algorithms under the current dataset")
            for key, name in alg_map.items():
                print(f"{key}: {name}")
            print("="*50)
            choice2 = input("Enter the algorithm number (or '0' for all) and press Enter: ").strip()
            if choice2 == '0':
                model_archs_to_process = valid_alg_archs
            elif choice2 in alg_map:
                model_archs_to_process = [alg_map[choice2]]
        else:
            logging.warning(f"No valid algorithm subfolders found in {os.path.basename(selected_dataset_dir)}.")
    else:
        logging.error("Invalid dataset selection. Exiting.")
        return

    if not model_archs_to_process or not selected_dataset_dir:
        return

    # 2. 确定输入尺寸
    if args.shape:
        analysis_shape = (args.shape[0], args.shape[1])
        logging.info(f"Using user-provided shape (H, W): {analysis_shape} for calculation.")
    else:
        analysis_shape = get_shape_from_path(selected_dataset_dir) or (config.IMAGE_HEIGHT, config.IMAGE_WIDTH)
        logging.info(f"Using detected/default shape (H, W): {analysis_shape} for calculation.")

    # 3. 主处理循环
    ori_analysis_shape = analysis_shape
    results: List[Dict[str, str]] = []
    for architecture in model_archs_to_process:
        logging.info(f"--- Processing model: {architecture} ---")
        
        # 特殊处理 DPT 模型的输入尺寸要求
        analysis_shape = ori_analysis_shape
        patch_size = 16  # From the encoder name '...patch16...'
        h, w = analysis_shape
        if h % patch_size != 0 or w % patch_size != 0:
                # Calculate the new, compatible dimensions by rounding up to the nearest multiple of patch_size
                new_h = ((h + patch_size - 1) // patch_size) * patch_size
                new_w = ((w + patch_size - 1) // patch_size) * patch_size
                analysis_shape = (new_h, new_w)
                logging.warning(f"DPT requires dimensions divisible by {patch_size}. "
                                f"Adjusting analysis shape from {ori_analysis_shape} to {analysis_shape}.")

        # 步骤 3.1: 查找最新的模型运行路径和权重文件路径
        latest_run_dir = find_latest_model_run_path(Path(selected_dataset_dir), architecture)
        checkpoint_path = None
        if latest_run_dir:
            # 权重文件路径，即用户期望的 best_model_save_path
            checkpoint_path = latest_run_dir / config.BEST_MODEL_SAVE_NAME

        # 步骤 3.2: 使用新函数加载模型（一行代码完成初始化和权重加载）
        model = load_model(architecture, checkpoint_path)
        
        if model is None:
            logging.warning(f"Skipping all analysis for {architecture} due to model initialization failure.")
            continue

        # 步骤 3.3: 在同一个模型实例上执行所有计算
        stats_flops_params = calculate_flops_params(model, analysis_shape)
        stats_fps = benchmark_model(model, analysis_shape)

        # 3.4 整合结果
        results.append({
            'Model': architecture,
            'Params': stats_flops_params['params'] if stats_flops_params else "N/A",
            'FLOPs': stats_flops_params['flops'] if stats_flops_params else "N/A",
            'Input_Shape (HxW)': f"{ori_analysis_shape[0]}x{ori_analysis_shape[1]}",
            'Average_FPS': f"{stats_fps['avg_fps']:.2f}" if stats_fps else "N/A",
            'FPS_Variance': f"{stats_fps['fps_variance']:.4f}" if stats_fps else "N/A"
        })

    # 4. 写入CSV文件
    if results:
        dataset_name = os.path.basename(selected_dataset_dir).replace('_outputs-SegModel', '')
        output_csv_path = os.path.join(selected_dataset_dir, f'{dataset_name}_flops_params_fps_summary.csv')
        
        try:
            with open(output_csv_path, 'w', newline='', encoding='utf-8') as csvfile:
                fieldnames = ['Model', 'Params', 'FLOPs', 'Input_Shape (HxW)', 'Average_FPS', 'FPS_Variance']
                writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
                writer.writeheader()
                writer.writerows(results)
            logging.info(f"=== All processing complete! Results saved to: {output_csv_path} ===")
        except IOError as e:
            logging.error(f"Could not write to CSV file: {output_csv_path}. Error: {e}")
    else:
        logging.info("No statistics were successfully generated, CSV file will not be created.")

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="SMP Model Automation and Evaluation Script")
    parser.add_argument(
        '--input_dir',
        type=str,
        default=str(config.PREDICT_ALL_BEST_MODELS_DIR),
        help="Root directory containing trained model folders."
    )
    parser.add_argument(
        '--shape',
        type=int,
        nargs=2,
        metavar=('HEIGHT', 'WIDTH'),
        help="Specify input shape (height width), overriding automatic detection."
    )
    args = parser.parse_args()
    main(args)