import os
import glob
import logging
import argparse
import re
import csv
from typing import Dict, Optional, List
from collections import defaultdict

# --- 新增导入：用于绘图 ---
try:
    import pandas as pd
    import matplotlib.pyplot as plt
    MATPLOTLIB_AVAILABLE = True
except ImportError:
    MATPLOTLIB_AVAILABLE = False
    logging.warning(
        "未找到 'pandas' 或 'matplotlib' 库。"
        "脚本将只生成CSV文件，无法自动绘图。"
        "请运行 'pip install pandas matplotlib' 来安装它们。"
    )
# --- 导入结束 ---


# --- 配置日志记录 ---
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# --- 辅助函数 (与 4_3 版本相同) ---

def find_all_log_files_sorted(algorithm_dir: str) -> List[str]:
    """
    查找给定算法目录中所有的.log文件，并按从新到旧的顺序排列。
    """
    try:
        subdirs = [d for d in os.listdir(algorithm_dir) if os.path.isdir(os.path.join(algorithm_dir, d))]
    except FileNotFoundError:
        logging.error(f"算法目录不存在: {algorithm_dir}")
        return []

    if not subdirs:
        logging.warning(f"在目录 {algorithm_dir} 中未找到任何时间戳子目录。")
        return []

    sorted_subdirs = sorted(subdirs, reverse=True)
    log_files = []
    for subdir_name in sorted_subdirs:
        subdir_path = os.path.join(algorithm_dir, subdir_name)
        logs_in_subdir = glob.glob(os.path.join(subdir_path, '*.log'))
        if logs_in_subdir:
            log_files.append(logs_in_subdir[0])
    
    return log_files

def get_max_epochs(config_path: str) -> Optional[int]:
    """
    从config.py文件中解析train_cfg字典以获取max_epochs的值。
    """
    if not os.path.exists(config_path):
        logging.error(f"配置文件不存在: {config_path}")
        return None
        
    try:
        with open(config_path, 'r', encoding='utf-8') as f:
            content = f.read()
        
        match = re.search(r"train_cfg\s*=\s*dict\(.*?max_epochs\s*=\s*(\d+),.*?\)", content, re.DOTALL)
        
        if match:
            max_epochs = int(match.group(1))
            logging.info(f"从 {os.path.basename(config_path)} 中成功读取 max_epochs: {max_epochs}")
            return max_epochs
        else:
            logging.warning(f"在 {config_path} 中未找到 'max_epochs'。")
            return None
    except Exception as e:
        logging.error(f"解析 {config_path} 时出错: {e}")
        return None

def parse_log_file_data(log_path: str, max_epochs: int) -> Dict:
    """
    解析日志文件，提取所有训练损失和所有验证mIoU。
    """
    try:
        with open(log_path, 'r', encoding='utf-8') as f:
            content = f.read()
    except IOError as e:
        logging.error(f"无法读取日志文件: {log_path}。错误: {e}")
        return {'training_losses': {}, 'validation_mious': []}

    training_losses_by_epoch = defaultdict(list)
    validation_mious = []

    # 1. 提取训练损失
    total_iters_match = re.search(r"Iter\(train\)\s*\[\s*(\d+)\s*/\s*(\d+)\]", content)
    total_iters = 0
    if total_iters_match:
        try:
            total_iters = int(total_iters_match.group(2))
        except ValueError:
            pass 
            
    if total_iters > 0:
        train_loss_pattern = re.compile(
            r"Iter\(train\)\s*\[\s*(\d+)\s*/\s*(\d+)\]"
            r"(?:.*?)loss:\s*([\d\.]+)"
        )
        for match in re.finditer(train_loss_pattern, content):
            try:
                current_iter = int(match.group(1))
                loss = float(match.group(3))
                epoch = int((current_iter / total_iters) * max_epochs)
                training_losses_by_epoch[epoch].append(loss)
            except (ValueError, ZeroDivisionError) as e:
                logging.warning(f"解析训练损失时出错: {e}")
    else:
         logging.warning(f"在 {os.path.basename(log_path)} 中未找到有效的 'Iter(train)' 行来确定总迭代次数。")

    # 2. 提取验证 mIoU
    val_summary_pattern = re.compile(
        r"Iter\(val\) \[\d+/\d+\]\s+aAcc:\s*[\d.]+\s+mIoU:\s*([\d.]+)\s+mAcc:\s*[\d.]+"
    )
    for match in re.finditer(val_summary_pattern, content):
        try:
            miou = float(match.group(1))
            validation_mious.append(miou)
        except ValueError:
            pass

    log_name = os.path.basename(log_path)
    if training_losses_by_epoch:
        logging.info(f"✅ 从 {log_name} 提取了 {len(training_losses_by_epoch)} 个Epoch的训练损失。")
    if validation_mious:
        logging.info(f"✅ 从 {log_name} 提取了 {len(validation_mious)} 次验证的mIoU。")
    if not training_losses_by_epoch and not validation_mious:
        logging.warning(f"❌ 在 {log_name} 中未找到训练损失或验证mIoU。")
        
    return {
        'training_losses': training_losses_by_epoch,
        'validation_mious': validation_mious
    }

# --- 新增：从 4_4 脚本中合并过来的绘图函数 ---

def plot_loss_curves(csv_path: str):
    """
    读取_training_loss_summary.csv文件并绘制训练损失曲线。
    (此版本已根据用户需求修改)

    Args:
        csv_path (str): 输入的CSV文件路径。
    """
    if not MATPLOTLIB_AVAILABLE:
        logging.warning("由于缺少 'pandas' 或 'matplotlib'，跳过绘图。")
        return

    if not os.path.exists(csv_path):
        logging.error(f"[绘图] 文件未找到: {csv_path}")
        return

    logging.info(f"[绘图] 正在读取数据: {os.path.basename(csv_path)}")
    try:
        df = pd.read_csv(csv_path)
    except Exception as e:
        logging.error(f"[绘图] 读取CSV时出错: {e}")
        return

    try:
        df = df.set_index('Algorithm')
    except KeyError:
        logging.error("[绘图] CSV文件中未找到 'Algorithm' 列。")
        return

    loss_cols = [col for col in df.columns if col.startswith('Epoch_') and col.endswith('_Loss')]
    
    if not loss_cols:
        logging.warning("[绘图] 在CSV中未找到任何 'Epoch_X_Loss' 列。")
        return

    try:
        epochs = [int(col.split('_')[1]) for col in loss_cols]
    except (ValueError, IndexError) as e:
        logging.error(f"[绘图] 解析Epoch列名时出错: {e}。列名格式应为 'Epoch_N_Loss'。")
        return
        
    df_losses = df[loss_cols].apply(pd.to_numeric, errors='coerce')

    # --- 开始绘图 ---
    logging.info("[绘图] 正在生成图表...")
    
    fig, ax = plt.subplots(figsize=(14, 8))

    for alg_name, row in df_losses.iterrows():
        # --- 修改点 1 ---
        # 移除 marker 和 markersize，使用实线 (linestyle='-')
        ax.plot(epochs, row.values, label=alg_name, linestyle='-')

    # --- 设置图表样式 ---
    ax.set_xlabel('Epoch', fontsize=12)
    ax.set_ylabel('Average Training Loss', fontsize=12)
    ax.set_title(f'Training Loss per Epoch\n(Source: {os.path.basename(csv_path)})', fontsize=14)
    ax.grid(True, linestyle=':', alpha=0.7)
    
    # --- 修改点 2 ---
    # 将Y轴(纵轴)的范围设置为 0 到 5
    ax.set_ylim(bottom=0, top=5) # TODO TODO
    
    # 将图例放在图表右侧外部
    ax.legend(bbox_to_anchor=(1.04, 1), loc='upper left', title="Algorithms")
    
    fig.tight_layout(rect=[0, 0, 0.85, 1]) 

    # --- 保存图表 ---
    output_png_path = os.path.splitext(csv_path)[0] + '.png'
    try:
        plt.savefig(output_png_path, dpi=150, bbox_inches='tight')
        logging.info(f"✅ 图表已成功保存到: {output_png_path}")
    except IOError as e:
        logging.error(f"[绘图] 保存图像时出错: {e}")
    finally:
        plt.close(fig) # 释放内存

# --- 主函数 (小幅修改以调用绘图) ---

def main(args):
    """
    脚本主入口，负责编排整个自动化分析流程。
    """
    input_root = args.input_dir
    output_root = args.output_dir
    
    if not os.path.isdir(input_root):
        logging.error(f"输入目录不存在: {input_root}")
        return

    # --- 交互式菜单 (保持不变) ---
    all_dataset_dirs = sorted([
        d for d in glob.glob(os.path.join(input_root, '*_outputs-MMSeg')) if os.path.isdir(d)
    ])
    if not all_dataset_dirs:
        logging.error(f"在输入目录 {input_root} 中未找到任何有效的数据集文件夹。")
        return
    dataset_map = {str(i + 1): path for i, path in enumerate(all_dataset_dirs)}
    print("\n" + "="*50)
    print("--- 步骤 1: 请选择要处理的数据集 ---")
    for key, path in dataset_map.items():
        print(f"{key}: {os.path.basename(path)}")
    print("="*50)
    choice1 = input("请输入数据集编号并按回车键: ").strip()

    model_dirs = []
    selected_dataset_dir = None
    
    if choice1 in dataset_map:
        selected_dataset_dir = dataset_map[choice1]
        logging.info(f"您已选择数据集: [{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"在 {os.path.basename(selected_dataset_dir)} 中未发现任何算法子文件夹。")
            return
        alg_map = {str(i + 1): path for i, path in enumerate(alg_dirs)}
        print("\n" + "="*50)
        print("--- 步骤 2: 请选择要处理的算法 ---")
        print("0: 批量处理当前数据集下的【全部】算法")
        for key, path in alg_map.items():
            print(f"{key}: {os.path.basename(path)}")
        print("="*50)
        choice2 = input("请输入算法编号 (或输入 '0' 处理全部) 并按回车键: ").strip()
        if choice2 == '0':
            model_dirs = alg_dirs
            logging.info(f"您选择了批量处理全部 {len(model_dirs)} 个算法。")
        elif choice2 in alg_map:
            model_dirs = [alg_map[choice2]]
            logging.info(f"您选择了处理单个算法: {os.path.basename(model_dirs[0])}")
        else:
            logging.error("无效的算法选择，程序已退出。")
            return
    else:
        logging.error("无效的数据集选择，程序已退出。")
        return
    # --- 交互式菜单结束 ---


    # --- 处理逻辑 (保持不变) ---
    csv_rows = []

    for model_dir in model_dirs:
        model_name = os.path.basename(model_dir)
        logging.info(f"\n--- 开始处理算法: {model_name} ---")

        all_logs_sorted = find_all_log_files_sorted(model_dir)
        
        if not all_logs_sorted:
            logging.warning(f"跳过算法 {model_name}，因为未找到任何日志文件。")
            continue
        
        latest_log_path = all_logs_sorted[0]
        config_path = os.path.join(os.path.dirname(latest_log_path), 'vis_data', 'config.py')
        max_epochs = get_max_epochs(config_path)
        if max_epochs is None:
            logging.error(f"无法为算法 {model_name} 找到 max_epochs，将跳过。")
            continue

        all_losses_for_model = defaultdict(list)
        all_mious_for_model = []

        for log_file_path in all_logs_sorted:
            logging.info(f"正在解析: {os.path.relpath(log_file_path)}")
            parsed_data = parse_log_file_data(log_file_path, max_epochs)
            for epoch, losses in parsed_data['training_losses'].items():
                all_losses_for_model[epoch].extend(losses)
            if parsed_data['validation_mious']:
                all_mious_for_model.extend(parsed_data['validation_mious'])

        if not all_losses_for_model and not all_mious_for_model:
            logging.warning(f"❌❌❌跳过算法 {model_name}，因为在其所有日志文件中都未能找到有效的训练或验证数据。❌❌❌")
            continue
            
        best_miou = 0.0
        if all_mious_for_model:
            try:
                best_miou = max(all_mious_for_model)
                logging.info(f"找到了最佳 mIoU: {best_miou:.4f}")
            except (ValueError, TypeError) as e:
                logging.error(f"为算法 {model_name} 寻找最佳mIoU时出错: {e}")
        else:
            logging.warning(f"算法 {model_name} 没有找到任何mIoU数据。")

        short_model_name = model_name.split('Alg_', 1)[1] if 'Alg_' in model_name else model_name
        row_data = {
            'Algorithm': short_model_name
        }
        
        sorted_epochs = sorted(all_losses_for_model.keys())
        if not sorted_epochs:
             logging.warning(f"算法 {model_name} 没有找到任何训练损失数据。")
             
        for epoch in sorted_epochs:
            losses = all_losses_for_model[epoch]
            avg_loss = sum(losses) / len(losses)
            row_data[f'Epoch_{epoch}_Loss'] = f"{avg_loss:.4f}"
        
        row_data['Best_mIoU'] = f"{best_miou:.4f}"
        csv_rows.append(row_data)
    # --- 处理逻辑结束 ---


    # --- 写入 CSV 文件 (修改后) ---
    if not csv_rows:
        logging.info("没有成功获取任何模型的统计数据，不生成 CSV 文件。")
        return

    # 动态生成所有列名
    all_fieldnames_set = set()
    for row in csv_rows:
        all_fieldnames_set.update(row.keys())

    base_fields = ['Algorithm']
    miou_field = ['Best_mIoU']
    loss_fields = [f for f in all_fieldnames_set if f.startswith('Epoch_')]
    
    try:
        loss_fields.sort(key=lambda x: int(x.split('_')[1]))
    except (ValueError, IndexError):
        logging.error("排序Epoch列名时出错，将按字母顺序排序。")
        loss_fields.sort() 

    final_fieldnames = base_fields + loss_fields + miou_field
    
    final_output_dir = os.path.join(output_root, os.path.basename(selected_dataset_dir))
    os.makedirs(final_output_dir, exist_ok=True)
    dataset_name = os.path.basename(selected_dataset_dir).split('_outputs-MMSeg')[0]
    output_csv_path = os.path.join(final_output_dir, f'{dataset_name}_training_loss_summary.csv')
    
    try:
        with open(output_csv_path, 'w', newline='', encoding='utf-8-sig') as csvfile:
            writer = csv.DictWriter(csvfile, fieldnames=final_fieldnames, extrasaction='ignore')
            writer.writeheader()
            writer.writerows(csv_rows)
        
        logging.info(f"\n=== CSV文件已成功保存到: {output_csv_path} ===")
        
        # --- 新增：调用绘图函数 ---
        # 仅在CSV成功写入后才尝试绘图
        plot_loss_curves(output_csv_path)
        # --- 新增结束 ---

    except IOError as e:
        logging.error(f"无法写入 CSV 文件: {output_csv_path}。错误: {e}")
    except Exception as e_plot:
        # 捕获绘图时可能发生的其他错误
        logging.error(f"在主流程中调用绘图时出错: {e_plot}")

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description="MMSegmentation 训练损失提取与绘图脚本 (V3-Integrated)"
    )
    parser.add_argument(
        '--input_dir',
        type=str,
        default='../Hardisk',
        help="包含数据集输出文件夹 (例如 '..._outputs-MMSeg') 的根目录。"
    )
    parser.add_argument(
        '--output_dir',
        type=str,
        default='../BestMode_Predict_Results_DataSet_Public',
        help="用于存储所有分析结果 (CSV和PNG) 的根目录。"
    )
    
    args = parser.parse_args()
    main(args)