utils/visualization.py

"""
可视化工具模块
包含训练过程可视化、模型解释性可视化等
"""

import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
import cv2
from PIL import Image
import torch
import torch.nn.functional as F
from typing import List, Tuple, Dict, Optional, Union
import pandas as pd
from matplotlib.animation import FuncAnimation
import warnings
warnings.filterwarnings('ignore')

# 设置matplotlib中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False


def plot_training_curves(train_history: Dict, save_path: str = None, 
                        figsize: Tuple[int, int] = (15, 10)):
    """
    绘制训练曲线
    
    Args:
        train_history: 训练历史字典
        save_path: 保存路径
        figsize: 图像大小
    """
    fig, axes = plt.subplots(2, 2, figsize=figsize)
    
    epochs = range(1, len(train_history['train_loss']) + 1)
    
    # 损失曲线
    axes[0, 0].plot(epochs, train_history['train_loss'], 'b-', label='训练损失', linewidth=2)
    axes[0, 0].plot(epochs, train_history['val_loss'], 'r-', label='验证损失', linewidth=2)
    axes[0, 0].set_title('损失曲线', fontsize=14, fontweight='bold')
    axes[0, 0].set_xlabel('Epoch')
    axes[0, 0].set_ylabel('Loss')
    axes[0, 0].legend()
    axes[0, 0].grid(True, alpha=0.3)
    
    # 准确率曲线
    axes[0, 1].plot(epochs, train_history['train_acc'], 'b-', label='训练准确率', linewidth=2)
    axes[0, 1].plot(epochs, train_history['val_acc'], 'r-', label='验证准确率', linewidth=2)
    axes[0, 1].set_title('准确率曲线', fontsize=14, fontweight='bold')
    axes[0, 1].set_xlabel('Epoch')
    axes[0, 1].set_ylabel('Accuracy (%)')
    axes[0, 1].legend()
    axes[0, 1].grid(True, alpha=0.3)
    
    # 学习率曲线
    if 'lr' in train_history:
        axes[1, 0].plot(epochs, train_history['lr'], 'g-', linewidth=2)
        axes[1, 0].set_title('学习率变化', fontsize=14, fontweight='bold')
        axes[1, 0].set_xlabel('Epoch')
        axes[1, 0].set_ylabel('Learning Rate')
        axes[1, 0].set_yscale('log')
        axes[1, 0].grid(True, alpha=0.3)
    
    # 训练摘要
    best_val_acc = max(train_history['val_acc'])
    best_epoch = train_history['val_acc'].index(best_val_acc) + 1
    final_train_loss = train_history['train_loss'][-1]
    final_val_loss = train_history['val_loss'][-1]
    
    summary_text = f"""训练摘要:
    
最佳验证准确率: {best_val_acc:.2f}%
最佳epoch: {best_epoch}
最终训练损失: {final_train_loss:.4f}
最终验证损失: {final_val_loss:.4f}
总训练轮数: {len(epochs)}
    """
    
    axes[1, 1].text(0.1, 0.9, summary_text, transform=axes[1, 1].transAxes, 
                    fontsize=12, verticalalignment='top',
                    bbox=dict(boxstyle='round', facecolor='lightblue', alpha=0.8))
    axes[1, 1].set_xlim(0, 1)
    axes[1, 1].set_ylim(0, 1)
    axes[1, 1].axis('off')
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"训练曲线已保存: {save_path}")
    
    return fig


def plot_class_distribution(data_dir: str, class_names: List[str] = None,
                           save_path: str = None, figsize: Tuple[int, int] = (12, 8)):
    """
    绘制数据集类别分布
    
    Args:
        data_dir: 数据目录
        class_names: 类别名称
        save_path: 保存路径
        figsize: 图像大小
    """
    import os
    
    class_counts = {}
    
    # 统计各类别样本数
    for class_idx, class_folder in enumerate(os.listdir(data_dir)):
        class_path = os.path.join(data_dir, class_folder)
        if os.path.isdir(class_path):
            count = len([f for f in os.listdir(class_path) 
                        if f.lower().endswith(('.png', '.jpg', '.jpeg'))])
            
            if class_names and class_idx < len(class_names):
                class_name = class_names[class_idx]
            else:
                class_name = class_folder
            
            class_counts[class_name] = count
    
    # 绘制条形图
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=figsize)
    
    # 条形图
    classes = list(class_counts.keys())
    counts = list(class_counts.values())
    colors = plt.cm.Set3(np.linspace(0, 1, len(classes)))
    
    bars = ax1.bar(classes, counts, color=colors, edgecolor='black', linewidth=1)
    ax1.set_title('类别样本分布', fontsize=14, fontweight='bold')
    ax1.set_xlabel('类别')
    ax1.set_ylabel('样本数量')
    ax1.tick_params(axis='x', rotation=45)
    
    # 添加数值标签
    for bar, count in zip(bars, counts):
        ax1.text(bar.get_x() + bar.get_width()/2, bar.get_height() + max(counts)*0.01,
                str(count), ha='center', va='bottom', fontweight='bold')
    
    # 饼图
    ax2.pie(counts, labels=classes, autopct='%1.1f%%', colors=colors,
           startangle=90, wedgeprops={'edgecolor': 'black'})
    ax2.set_title('类别比例分布', fontsize=14, fontweight='bold')
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"类别分布图已保存: {save_path}")
    
    return fig, class_counts


def plot_sample_images(data_dir: str, class_names: List[str], 
                      samples_per_class: int = 3,
                      save_path: str = None, figsize: Tuple[int, int] = (15, 10)):
    """
    显示每个类别的样本图像
    
    Args:
        data_dir: 数据目录
        class_names: 类别名称
        samples_per_class: 每类显示的样本数
        save_path: 保存路径
        figsize: 图像大小
    """
    import os
    import random
    
    n_classes = len(class_names)
    fig, axes = plt.subplots(n_classes, samples_per_class, figsize=figsize)
    
    if n_classes == 1:
        axes = [axes]
    
    for class_idx, class_name in enumerate(class_names):
        class_dir = os.path.join(data_dir, class_name)
        if not os.path.exists(class_dir):
            # 尝试用索引查找目录
            class_dirs = [d for d in os.listdir(data_dir) if os.path.isdir(os.path.join(data_dir, d))]
            if class_idx < len(class_dirs):
                class_dir = os.path.join(data_dir, class_dirs[class_idx])
        
        if os.path.exists(class_dir):
            # 获取图像文件
            image_files = [f for f in os.listdir(class_dir) 
                          if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
            
            # 随机选择样本
            selected_files = random.sample(image_files, 
                                         min(samples_per_class, len(image_files)))
            
            for sample_idx, img_file in enumerate(selected_files):
                img_path = os.path.join(class_dir, img_file)
                img = cv2.imread(img_path)
                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
                
                if samples_per_class == 1:
                    ax = axes[class_idx]
                else:
                    ax = axes[class_idx, sample_idx]
                
                ax.imshow(img)
                ax.axis('off')
                
                if sample_idx == 0:
                    ax.set_ylabel(class_name, fontsize=12, fontweight='bold')
            
            # 填充空白位置
            for sample_idx in range(len(selected_files), samples_per_class):
                if samples_per_class == 1:
                    ax = axes[class_idx]
                else:
                    ax = axes[class_idx, sample_idx]
                ax.axis('off')
    
    plt.suptitle('各类别样本展示', fontsize=16, fontweight='bold')
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"样本图像已保存: {save_path}")
    
    return fig


def plot_model_comparison(results: List[Dict], save_path: str = None,
                         figsize: Tuple[int, int] = (15, 10)):
    """
    比较不同模型的性能
    
    Args:
        results: 模型结果列表，每个元素包含模型名称和指标
        save_path: 保存路径
        figsize: 图像大小
    """
    metrics_to_plot = ['accuracy', 'macro_f1', 'macro_precision', 'macro_recall']
    model_names = [r['model_name'] for r in results]
    
    fig, axes = plt.subplots(2, 2, figsize=figsize)
    axes = axes.ravel()
    
    for idx, metric in enumerate(metrics_to_plot):
        values = [r['metrics'].get(metric, 0) for r in results]
        colors = plt.cm.Set2(np.linspace(0, 1, len(model_names)))
        
        bars = axes[idx].bar(model_names, values, color=colors, 
                           edgecolor='black', linewidth=1)
        
        axes[idx].set_title(f'{metric.replace("_", " ").title()}', 
                          fontsize=12, fontweight='bold')
        axes[idx].set_ylabel('Score')
        axes[idx].tick_params(axis='x', rotation=45)
        axes[idx].grid(True, alpha=0.3)
        
        # 添加数值标签
        for bar, value in zip(bars, values):
            axes[idx].text(bar.get_x() + bar.get_width()/2, 
                         bar.get_height() + max(values)*0.01,
                         f'{value:.3f}', ha='center', va='bottom', fontweight='bold')
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"模型比较图已保存: {save_path}")
    
    return fig


def create_interactive_training_dashboard(train_history: Dict, save_path: str = None):
    """
    创建交互式训练仪表板
    
    Args:
        train_history: 训练历史
        save_path: 保存路径（HTML文件）
    """
    # 惰性导入 plotly，未安装则给出提示
    try:
        import plotly.graph_objects as go
        from plotly.subplots import make_subplots
    except Exception as exc:
        raise ImportError(
            "使用 create_interactive_training_dashboard 需要安装 plotly，请运行: pip install plotly"
        ) from exc

    epochs = list(range(1, len(train_history['train_loss']) + 1))
    
    # 创建子图
    fig = make_subplots(
        rows=2, cols=2,
        subplot_titles=('损失曲线', '准确率曲线', '学习率变化', '训练摘要'),
        specs=[[{"secondary_y": False}, {"secondary_y": False}],
               [{"secondary_y": False}, {"type": "table"}]]
    )
    
    # 损失曲线
    fig.add_trace(
        go.Scatter(x=epochs, y=train_history['train_loss'], 
                  name='训练损失', line=dict(color='blue')),
        row=1, col=1
    )
    fig.add_trace(
        go.Scatter(x=epochs, y=train_history['val_loss'], 
                  name='验证损失', line=dict(color='red')),
        row=1, col=1
    )
    
    # 准确率曲线
    fig.add_trace(
        go.Scatter(x=epochs, y=train_history['train_acc'], 
                  name='训练准确率', line=dict(color='blue')),
        row=1, col=2
    )
    fig.add_trace(
        go.Scatter(x=epochs, y=train_history['val_acc'], 
                  name='验证准确率', line=dict(color='red')),
        row=1, col=2
    )
    
    # 学习率变化
    if 'lr' in train_history:
        fig.add_trace(
            go.Scatter(x=epochs, y=train_history['lr'], 
                      name='学习率', line=dict(color='green')),
            row=2, col=1
        )
        fig.update_yaxes(type="log", row=2, col=1)
    
    # 训练摘要表格
    best_val_acc = max(train_history['val_acc'])
    best_epoch = train_history['val_acc'].index(best_val_acc) + 1
    
    summary_data = [
        ['指标', '数值'],
        ['最佳验证准确率', f'{best_val_acc:.2f}%'],
        ['最佳Epoch', str(best_epoch)],
        ['最终训练损失', f'{train_history["train_loss"][-1]:.4f}'],
        ['最终验证损失', f'{train_history["val_loss"][-1]:.4f}'],
        ['总训练轮数', str(len(epochs))]
    ]
    
    fig.add_trace(
        go.Table(
            header=dict(values=summary_data[0], fill_color='lightblue'),
            cells=dict(values=list(zip(*summary_data[1:])), fill_color='white')
        ),
        row=2, col=2
    )
    
    # 更新布局
    fig.update_layout(
        title='训练过程可视化仪表板',
        height=800,
        showlegend=False
    )
    
    if save_path:
        fig.write_html(save_path)
        print(f"交互式仪表板已保存: {save_path}")
    
    return fig


def visualize_feature_maps(model: torch.nn.Module, image: torch.Tensor, 
                          layer_name: str, save_path: str = None,
                          figsize: Tuple[int, int] = (20, 15)):
    """
    可视化特征图
    
    Args:
        model: PyTorch模型
        image: 输入图像tensor
        layer_name: 要可视化的层名称
        save_path: 保存路径
        figsize: 图像大小
    """
    # 注册hook获取特征图
    features = {}
    
    def hook_fn(module, input, output):
        features['feature_map'] = output
    
    # 找到目标层并注册hook
    target_layer = None
    for name, module in model.named_modules():
        if layer_name in name:
            target_layer = module
            break
    
    if target_layer is None:
        print(f"未找到层: {layer_name}")
        return None
    
    handle = target_layer.register_forward_hook(hook_fn)
    
    # 前向传播
    model.eval()
    with torch.no_grad():
        _ = model(image.unsqueeze(0))
    
    # 移除hook
    handle.remove()
    
    # 获取特征图
    feature_map = features['feature_map'].squeeze(0)  # 移除batch维度
    n_features = feature_map.shape[0]
    
    # 计算网格大小
    n_cols = 8
    n_rows = (n_features + n_cols - 1) // n_cols
    
    fig, axes = plt.subplots(n_rows, n_cols, figsize=figsize)
    if n_rows == 1:
        axes = [axes]
    axes = np.array(axes).ravel()
    
    for i in range(n_features):
        feature = feature_map[i].cpu().numpy()
        
        # 标准化到[0,1]
        feature = (feature - feature.min()) / (feature.max() - feature.min() + 1e-8)
        
        axes[i].imshow(feature, cmap='viridis')
        axes[i].set_title(f'Feature {i+1}')
        axes[i].axis('off')
    
    # 隐藏多余的子图
    for i in range(n_features, len(axes)):
        axes[i].axis('off')
    
    plt.suptitle(f'特征图可视化 - {layer_name}', fontsize=16, fontweight='bold')
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"特征图已保存: {save_path}")
    
    return fig


def plot_attention_weights(attention_weights: np.ndarray, tokens: List[str] = None,
                          save_path: str = None, figsize: Tuple[int, int] = (12, 10)):
    """
    可视化注意力权重（适用于Vision Transformer）
    
    Args:
        attention_weights: 注意力权重矩阵
        tokens: token标签
        save_path: 保存路径
        figsize: 图像大小
    """
    plt.figure(figsize=figsize)
    
    if tokens is None:
        tokens = [f'Token {i+1}' for i in range(attention_weights.shape[0])]
    
    # 创建热力图
    sns.heatmap(attention_weights, 
                xticklabels=tokens, yticklabels=tokens,
                cmap='Blues', annot=False, square=True,
                cbar_kws={'label': 'Attention Weight'})
    
    plt.title('注意力权重可视化', fontsize=16, fontweight='bold')
    plt.xlabel('Key Tokens')
    plt.ylabel('Query Tokens')
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"注意力权重图已保存: {save_path}")
    
    return plt.gcf()


def create_prediction_gallery(images: List[str], predictions: List[Dict],
                             true_labels: List[str] = None, 
                             save_path: str = None,
                             images_per_row: int = 4,
                             figsize: Tuple[int, int] = (20, 15)):
    """
    创建预测结果画廊
    
    Args:
        images: 图像路径列表
        predictions: 预测结果列表
        true_labels: 真实标签列表
        save_path: 保存路径
        images_per_row: 每行图像数
        figsize: 图像大小
    """
    n_images = len(images)
    n_rows = (n_images + images_per_row - 1) // images_per_row
    
    fig, axes = plt.subplots(n_rows, images_per_row, figsize=figsize)
    if n_rows == 1:
        axes = [axes]
    axes = np.array(axes).ravel()
    
    for i, (img_path, pred) in enumerate(zip(images, predictions)):
        if i >= len(axes):
            break
            
        # 读取并显示图像
        img = cv2.imread(img_path)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        axes[i].imshow(img)
        
        # 创建标题
        title = f"预测: {pred['predicted_label']}\n"
        title += f"置信度: {pred['confidence']:.3f}"
        
        if true_labels and i < len(true_labels):
            title = f"真实: {true_labels[i]}\n" + title
            # 如果预测错误，用红色标题
            if pred['predicted_label'] != true_labels[i]:
                axes[i].set_title(title, color='red', fontweight='bold')
            else:
                axes[i].set_title(title, color='green', fontweight='bold')
        else:
            axes[i].set_title(title, fontweight='bold')
        
        axes[i].axis('off')
    
    # 隐藏多余的子图
    for i in range(len(images), len(axes)):
        axes[i].axis('off')
    
    plt.suptitle('预测结果画廊', fontsize=16, fontweight='bold')
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"预测画廊已保存: {save_path}")
    
    return fig


if __name__ == "__main__":
    # 测试可视化功能
    
    # 模拟训练历史数据
    epochs = 50
    train_history = {
        'train_loss': [1.5 * np.exp(-0.1 * i) + 0.1 + 0.05 * np.random.randn() for i in range(epochs)],
        'val_loss': [1.6 * np.exp(-0.08 * i) + 0.15 + 0.08 * np.random.randn() for i in range(epochs)],
        'train_acc': [60 * (1 - np.exp(-0.1 * i)) + 10 * np.random.randn() for i in range(epochs)],
        'val_acc': [55 * (1 - np.exp(-0.08 * i)) + 12 * np.random.randn() for i in range(epochs)],
        'lr': [0.001 * (0.9 ** (i // 10)) for i in range(epochs)]
    }
    
    # 绘制训练曲线
    fig = plot_training_curves(train_history, 'test_training_curves.png')
    plt.close()
    
    # 创建交互式仪表板
    interactive_fig = create_interactive_training_dashboard(train_history, 'test_dashboard.html')
    
    print("可视化测试完成！")