import os
import glob
import logging
import argparse
import re
import csv
import numpy as np
from typing import Dict, Optional, List

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

# --- Helper Functions (No changes in this section) ---

def find_smp_csv_file(algorithm_dir: str) -> Optional[str]:
    """
    Finds the 'training_metrics.csv' file in a given algorithm directory.
    """
    csv_path = os.path.join(algorithm_dir, 'training_metrics.csv')
    if os.path.exists(csv_path):
        return csv_path
    else:
        logging.warning(f"Could not find 'training_metrics.csv' in {algorithm_dir}")
        return None

def parse_smp_metrics(csv_path: str) -> Optional[Dict]:
    """
    Parses 'training_metrics.csv' to find the best epoch based on the highest
    IoU score and extracts all relevant metrics.
    """
    try:
        with open(csv_path, 'r', encoding='utf-8') as f:
            reader = csv.DictReader(f)
            rows = list(reader)
    except (IOError, StopIteration) as e:
        logging.error(f"Cannot read or parse CSV file: {csv_path}. Error: {e}")
        return None

    if not rows:
        logging.warning(f"❌ CSV file is empty: {os.path.basename(csv_path)}")
        return None

    try:
        best_row = max(rows, key=lambda row: float(row.get('iou_score', 0)))
        best_iou = float(best_row.get('iou_score', 0))
        if best_iou == 0:
             logging.warning(f"❌ Best 'iou_score' is 0 in {os.path.basename(csv_path)}. Check training.")
             return None
    except (ValueError, TypeError) as e:
        logging.error(f"Could not find a valid 'iou_score' in {csv_path}. Error: {e}")
        return None

    results = {
        'epoch': best_row.get('epoch', 'N/A'),
        'summary': {
            'mIoU': best_row.get('iou_score', 'N/A'),
            'mAcc': 'N/A',
            'aAcc': best_row.get('accuracy', 'N/A')
        },
        'class_wise': []
    }

    per_class_accuracies = []
    num_classes = 0
    class_keys = [key for key in best_row if key.startswith('tp_class_')]
    if class_keys:
        num_classes = max([int(k.split('_')[-1]) for k in class_keys]) + 1

    for i in range(num_classes):
        try:
            tp = float(best_row.get(f'tp_class_{i}', 0))
            fp = float(best_row.get(f'fp_class_{i}', 0))
            fn = float(best_row.get(f'fn_class_{i}', 0))
            
            class_iou = tp / (tp + fp + fn + 1e-6)
            class_acc_str = best_row.get(f'cpa_class_{i}', '0')
            per_class_accuracies.append(float(class_acc_str))
            
            class_name = '背景' if str(i) == '0' else str(i)
            
            results['class_wise'].append({
                'Class': class_name,
                'IoU': f"{class_iou:.4f}",
                'Acc': class_acc_str
            })
        except (ValueError, TypeError) as e:
            logging.warning(f"Could not process metrics for class {i}. Error: {e}")
            continue
    
    if per_class_accuracies:
        results['summary']['mAcc'] = f"{np.mean(per_class_accuracies):.4f}"

    if results['class_wise']:
        logging.info(f"✅ Successfully parsed metrics for epoch '{results['epoch']}' from {os.path.basename(csv_path)}.")
        return results
    else:
        logging.warning(f"❌ Could not parse any per-class metrics from {os.path.basename(csv_path)}.")
        return None


# --- Main Function ---
def main(args):
    """
    Main script entry point to orchestrate the analysis workflow.
    """
    input_root = args.input_dir
    output_root = args.output_dir
    
    if not os.path.isdir(input_root):
        logging.error(f"Input directory does not exist: {input_root}")
        return

    # --- Interactive Menu ---
    all_dataset_dirs = sorted([
        d for d in glob.glob(os.path.join(input_root, '*_outputs-SegModel')) if os.path.isdir(d)
    ])

    if not all_dataset_dirs:
        logging.error(f"No valid dataset folders ending in '_outputs-SegModel' found in {input_root}.")
        return

    dataset_map = {str(i + 1): path for i, path in enumerate(all_dataset_dirs)}
    
    print("\n" + "="*50)
    print("--- Step 1: Please select a 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_dirs = []
    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 = sorted([
            d for d in glob.glob(os.path.join(selected_dataset_dir, '*')) if os.path.isdir(d)
        ])

        if not alg_dirs:
            logging.warning(f"No algorithm subfolders found in {os.path.basename(selected_dataset_dir)}.")
            return
            
        alg_map = {str(i + 1): path for i, path in enumerate(alg_dirs)}
        print("\n" + "="*50)
        print("--- Step 2: Please select the algorithm(s) to process ---")
        print("0: Process [ALL] algorithms under the current dataset")
        for key, path in alg_map.items():
            print(f"{key}: {os.path.basename(path)}")
        print("="*50)

        choice2 = input("Enter the algorithm number (or '0' for all) and press Enter: ").strip()

        if choice2 == '0':
            model_dirs = alg_dirs
            logging.info(f"You have chosen to batch process all {len(model_dirs)} algorithms.")
        elif choice2 in alg_map:
            model_dirs = [alg_map[choice2]]
            logging.info(f"You have chosen to process a single algorithm: {os.path.basename(model_dirs[0])}")
        else:
            logging.error("Invalid algorithm selection. The program will exit.")
            return
    else:
        logging.error("Invalid dataset selection. The program will exit.")
        return

    # --- Start processing the selected algorithms ---
    csv_rows = []

    for model_dir in model_dirs:
        model_name_full = os.path.basename(model_dir)
        logging.info(f"\n--- Processing algorithm: {model_name_full} ---")

        csv_file_path = find_smp_csv_file(model_dir)
        
        if not csv_file_path:
            logging.warning(f"Skipping {model_name_full} as no 'training_metrics.csv' was found.")
            continue
        
        metrics = parse_smp_metrics(csv_file_path)
        
        if not metrics:
            logging.warning(f"❌❌❌ Skipping {model_name_full} due to a failure in parsing metrics. ❌❌❌")
            continue
            
        summary = metrics['summary']
        short_model_name = model_name_full.split('_')[0]
        
        row_data = {
            'Algorithm': short_model_name,
            'Epoch': metrics['epoch'],
            'mIoU': f"{round(float(summary['mIoU']), 4)*100:.2f}",
            'mAcc': f"{round(float(summary['mAcc']), 4)*100:.2f}",
            'aAcc': f"{round(float(summary['aAcc']), 4)*100:.2f}"
        }
        
        for class_data in metrics['class_wise']:
            class_name = class_data['Class']
            row_data[f'{class_name}_IoU'] = f"{round(float(class_data['IoU']), 4)*100:.2f}"
            row_data[f'{class_name}_Acc'] = f"{round(float(class_data['Acc']), 4)*100:.2f}"
        
        csv_rows.append(row_data)

    # --- Write results to the final CSV file ---
    if not csv_rows:
        logging.info("No data was successfully collected. No CSV file will be generated.")
        return

    # --- NEW: DYNAMICALLY GENERATE FIELDNAMES ---
    base_fieldnames = ['Algorithm', 'Epoch', 'mIoU', 'mAcc', 'aAcc']
    
    # Collect all unique per-class fieldnames from the data
    extra_fieldnames = set()
    for row in csv_rows:
        for key in row.keys():
            if key not in base_fieldnames:
                extra_fieldnames.add(key)
    
    # Define a sort key for natural sorting (e.g., '1', '2', '10' instead of '1', '10', '2')
    def natural_sort_key(key_name):
        parts = key_name.split('_')
        class_part = parts[0]
        metric_part = parts[1] if len(parts) > 1 else ''
        if class_part == '背景':
            return (-1, metric_part)  # Put '背景' (background) first
        try:
            return (int(class_part), metric_part)
        except ValueError:
            return (float('inf'), key_name) # Put any non-numeric names last
    
    # Combine base fields with sorted extra fields
    final_fieldnames = base_fieldnames + sorted(list(extra_fieldnames), key=natural_sort_key)
    
    # Build the output path
    output_dataset_folder = os.path.basename(selected_dataset_dir)
    final_output_dir = os.path.join(output_root, output_dataset_folder)
    os.makedirs(final_output_dir, exist_ok=True)
    
    dataset_name = output_dataset_folder.replace('_outputs-SegModel', '')
    output_csv_path = os.path.join(final_output_dir, f'{dataset_name}_metrics_summary_wide.csv')
    
    try:
        with open(output_csv_path, 'w', newline='', encoding='utf-8-sig') as csvfile:
            writer = csv.DictWriter(csvfile, fieldnames=final_fieldnames)
            writer.writeheader()
            writer.writerows(csv_rows)
        
        logging.info(f"\n=== All processing is complete! Results have been saved to: {output_csv_path} ===")
    except IOError as e:
        logging.error(f"Failed to write to CSV file: {output_csv_path}. Error: {e}")

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="SMP (Segmentation Models Pytorch) Final Metrics Extraction Script")
    parser.add_argument(
        '--input_dir',
        type=str,
        default='../Hardisk',
        help="The root directory containing dataset output folders (e.g., '..._outputs-SegModel')."
    )
    parser.add_argument(
        '--output_dir',
        type=str,
        default='../BestMode_Predict_Results_DataSet_Public',
        help="The root directory where analysis results will be stored."
    )
    
    args = parser.parse_args()
    main(args)