import os
import glob
import logging
import argparse
import csv
from typing import Dict, Optional, Tuple, List

import torch
from fvcore.nn import FlopCountAnalysis

from train import initialize_components
import config  # Assumes config.py is in the same directory or accessible

import time
import numpy as np
from pathlib import Path
import sys

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.py --input_dir '../BestMode_Predict_Results_DataSet_Public'
# python 2_predict_params_and_FLOPs.py --shape 512 512 # --shape 1080 1920

def get_flops_and_params_smp(architecture: str, shape: Tuple[int, int]) -> Optional[Dict[str, str]]:
    """
    Calculates FLOPs and parameters for a given SMP model architecture using fvcore.

    Args:
        architecture (str): The name of the model architecture.
        shape (Tuple[int, int]): The input image shape (H, W).

    Returns:
        Optional[Dict[str, str]]: A dictionary with 'params' and 'flops', or None on failure.
    """
    logging.info(f"Analyzing model: '{architecture}' with input shape {shape}...")

    num_classes = len(config.CLASSES)
    model, _, _, _ = initialize_components(architecture, num_classes, config.SEG_MODE)

    if model is None:
        return None

    model.to(config.DEVICE)
    model.eval()

    dummy_input = torch.randn(1, 3, shape[0], shape[1]).to(config.DEVICE)

    try:
        # Calculate FLOPs using fvcore
        flops_analyzer = FlopCountAnalysis(model, dummy_input)
        total_flops = flops_analyzer.total()
        gflops = total_flops / 1e9

        # Calculate Parameters
        total_params = sum(p.numel() for p in model.parameters())
        m_params = total_params / 1e6

        logging.info(f"✅ Analysis successful for '{architecture}': {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 for '{architecture}'. Error: {e}")
        return None

def get_shape_from_path(path: str) -> Optional[Tuple[int, int]]:
    """
    Extracts resolution (WxH) from a directory path using regex.

    Args:
        path (str): The path to the dataset directory.

    Returns:
        Optional[Tuple[int, int]]: A tuple of (Height, Width), or None if not found.
    """
    # 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)  # Return as H, W
    else:
        return None

# --- 主函数 ---
def main(args):
    """
    Main script entry point for the automated analysis workflow.
    """
    input_root = args.input_dir

    if not os.path.isdir(input_root):
        logging.error(f"Input directory does not exist: {input_root}")
        return

    # 1. Find all valid dataset directories (ending with '-SegModel')
    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 (e.g., '*_outputs-SegModel') found in {input_root}.")
        return

    # 2. First-level menu: Select dataset
    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)}]")

        # 3. Find all algorithm subdirectories and extract base architecture names
        alg_dirs_full_path = sorted([d for d in glob.glob(os.path.join(selected_dataset_dir, '*')) if os.path.isdir(d)])
        
        # CORRECTED LOGIC: Create a list of valid base architecture names found in the directory
        valid_alg_archs = []
        for path in alg_dirs_full_path:
            # Extract base name, e.g., 'Unet' from 'Unet_2025-09-24_14-24-05'
            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 not valid_alg_archs:
            logging.warning(f"No valid algorithm subfolders found in {os.path.basename(selected_dataset_dir)}. Please check folder names match keys in config.ALL_MODEL_CONFIGS.")
        else:
            # 4. Second-level menu: Select algorithm using the base names
            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()

            # 5. Determine the final list of architectures to process
            if choice2 == '0':
                model_archs_to_process = valid_alg_archs
                logging.info(f"You have chosen to process all {len(model_archs_to_process)} algorithms.")
            elif choice2 in alg_map:
                model_archs_to_process = [alg_map[choice2]]
                logging.info(f"You have chosen to process a single algorithm: {model_archs_to_process[0]}")
            else:
                logging.error("Invalid algorithm selection. Exiting.")
    else:
        logging.error("Invalid dataset selection. Exiting.")

    if not model_archs_to_process or not selected_dataset_dir:
        return

    # --- Processing logic ---
    results: List[Dict[str, str]] = []
    
    # Determine input shape for analysis
    input_shape = None
    if args.shape:
        # 1. Prioritize user-provided shape from command line
        input_shape = (args.shape[0], args.shape[1])
        logging.info(f"Using user-provided shape (H, W): {input_shape} for calculation.")
    else:
        # 2. Fallback to detecting shape from folder name
        logging.info("No --shape argument provided. Attempting to detect from folder name...")
        input_shape = get_shape_from_path(selected_dataset_dir)
        if not input_shape:
            # 3. Fallback to config.py defaults if detection fails
            logging.warning(f"Could not automatically detect resolution from folder '{os.path.basename(selected_dataset_dir)}'.")
            logging.info(f"Using default shape from config.py: ({config.IMAGE_HEIGHT}, {config.IMAGE_WIDTH})")
            input_shape = (config.IMAGE_HEIGHT, config.IMAGE_WIDTH)
        else:
            logging.info(f"Detected input shape (H, W): {input_shape} from folder name.")

    for architecture in model_archs_to_process:
        logging.info(f"--- Processing model: {architecture} ---")
        
        # ==================== SOLUTION CODE BLOCK START ====================
        # For Transformer models like DPT, ensure input shape is divisible by the patch size.
        analysis_shape = input_shape
        if architecture == 'DPT':
            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 {input_shape} to {analysis_shape}.")
        # ===================== SOLUTION CODE BLOCK END =====================

        stats = get_flops_and_params_smp(architecture, analysis_shape)
        if stats:
            results.append({
                'Model': architecture,
                'Params': stats['params'],
                'FLOPs': stats['flops'],
                'Input_Shape (HxW)': f"{analysis_shape[0]}x{analysis_shape[1]}"
            })
        else:
            logging.warning(f"Failed to get stats for model {architecture}.")

    # --- Write results to CSV file ---
    if not results:
        logging.info("No statistics were successfully generated, CSV file will not be created.")
        return

    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_summary.csv')
    
    try:
        with open(output_csv_path, 'w', newline='', encoding='utf-8') as csvfile:
            fieldnames = ['Model', 'Params', 'FLOPs', 'Input_Shape (HxW)']
            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}")

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 (e.g., 'DATASET_outputs-SegModel')."
    )
    parser.add_argument(
        '--shape',
        type=int,
        nargs=2,
        metavar=('HEIGHT', 'WIDTH'),
        help="Specify input shape (height width) for analysis, overriding automatic detection."
    )
    args = parser.parse_args()
    main(args)