Direction9/experiments/visualization.py

"""
实验可视化模块

功能:
- 绘制谱距离热力图
- 绘制有效秩柱状图
- 绘制谱分布对比图
"""

import os
import numpy as np
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Any
import json

# 延迟导入 matplotlib，避免不必要的依赖
_mpl_available = False
try:
    import matplotlib
    matplotlib.use('Agg')  # 非交互式后端
    import matplotlib.pyplot as plt
    import seaborn as sns
    _mpl_available = True
except ImportError:
    pass


def check_mpl_available():
    """检查 matplotlib 是否可用"""
    if not _mpl_available:
        raise ImportError(
            "matplotlib 或 seaborn 未安装。请运行：pip install matplotlib seaborn"
        )


def plot_heatmap(
    matrix: np.ndarray,
    labels: List[str],
    title: str = "Spectrum Distance Matrix",
    cmap: str = "viridis",
    output_path: Optional[str] = None,
    figsize: Tuple[int, int] = (8, 6),
) -> Optional[bytes]:
    """
    绘制热力图

    Args:
        matrix: n×n 距离矩阵
        labels: 模型/领域标签
        title: 图标题
        cmap: 颜色映射
        output_path: 输出路径（None 则返回 bytes）
        figsize: 图像大小

    Returns:
        图像 bytes 或 None（如果保存到文件）
    """
    check_mpl_available()

    fig, ax = plt.subplots(figsize=figsize)

    sns.heatmap(
        matrix,
        annot=True,
        fmt=".3f",
        cmap=cmap,
        xticklabels=labels,
        yticklabels=labels,
        ax=ax,
        cbar_kws={"label": "Wasserstein-1 Distance"},
    )

    ax.set_title(title)
    ax.set_xlabel("Model")
    ax.set_ylabel("Model")

    plt.tight_layout()

    if output_path:
        Path(output_path).parent.mkdir(parents=True, exist_ok=True)
        plt.savefig(output_path, dpi=150, bbox_inches="tight")
        plt.close(fig)
        return None
    else:
        from io import BytesIO
        buf = BytesIO()
        plt.savefig(buf, format="png", dpi=150, bbox_inches="tight")
        plt.close(fig)
        buf.seek(0)
        return buf.read()


def plot_bar_chart(
    data: Dict[str, float],
    title: str = "Effective Rank Comparison",
    xlabel: str = "Model",
    ylabel: str = "Effective Rank",
    output_path: Optional[str] = None,
    figsize: Tuple[int, int] = (10, 6),
    color: str = "steelblue",
    errorBars: Optional[Dict[str, float]] = None,
) -> Optional[bytes]:
    """
    绘制柱状图

    Args:
        data: {label: value, ...}
        title: 图标题
        xlabel: X 轴标签
        ylabel: Y 轴标签
        output_path: 输出路径
        figsize: 图像大小
        color: 柱子颜色
        errorBars: {label: std, ...} 误差线

    Returns:
        图像 bytes 或 None
    """
    check_mpl_available()

    labels = list(data.keys())
    values = list(data.values())

    fig, ax = plt.subplots(figsize=figsize)

    x_pos = range(len(labels))

    if errorBars:
        errors = [errorBars.get(l, 0) for l in labels]
        ax.bar(x_pos, values, yerr=errors, capsize=5, color=color, alpha=0.8)
    else:
        ax.bar(x_pos, values, color=color, alpha=0.8)

    ax.set_xticks(x_pos)
    ax.set_xticklabels(labels, rotation=45, ha="right")
    ax.set_title(title)
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)

    plt.tight_layout()

    if output_path:
        Path(output_path).parent.mkdir(parents=True, exist_ok=True)
        plt.savefig(output_path, dpi=150, bbox_inches="tight")
        plt.close(fig)
        return None
    else:
        from io import BytesIO
        buf = BytesIO()
        plt.savefig(buf, format="png", dpi=150, bbox_inches="tight")
        plt.close(fig)
        buf.seek(0)
        return buf.read()


def plot_spectrum_comparison(
    prob_dict: Dict[str, np.ndarray],
    title: str = "Spectrum Distribution Comparison",
    output_path: Optional[str] = None,
    figsize: Tuple[int, int] = (10, 6),
    log_scale: bool = True,
) -> Optional[bytes]:
    """
    绘制多个模型的谱分布对比图

    Args:
        prob_dict: {model_name: prob_distribution, ...}
        title: 图标题
        output_path: 输出路径
        figsize: 图像大小
        log_scale: Y 轴是否使用对数刻度

    Returns:
        图像 bytes 或 None
    """
    check_mpl_available()

    fig, ax = plt.subplots(figsize=figsize)

    colors = plt.cm.tab10.colors
    markers = ["o", "s", "^", "D", "v", "<", ">", "p"]

    for i, (model_name, prob) in enumerate(prob_dict.items()):
        color = colors[i % len(colors)]
        marker = markers[i % len(markers)]

        x = range(len(prob))
        ax.plot(
            x, prob,
            marker=marker,
            markersize=4,
            linestyle="-",
            label=model_name,
            color=color,
            alpha=0.8,
        )

    ax.set_title(title)
    ax.set_xlabel("Eigenvalue Index")
    ax.set_ylabel("Normalized Probability")

    if log_scale:
        ax.set_yscale("log")

    ax.legend(loc="upper right", fontsize=8)
    ax.grid(True, alpha=0.3)

    plt.tight_layout()

    if output_path:
        Path(output_path).parent.mkdir(parents=True, exist_ok=True)
        plt.savefig(output_path, dpi=150, bbox_inches="tight")
        plt.close(fig)
        return None
    else:
        from io import BytesIO
        buf = BytesIO()
        plt.savefig(buf, format="png", dpi=150, bbox_inches="tight")
        plt.close(fig)
        buf.seek(0)
        return buf.read()


def plot_convergence_analysis(
    result_dict: Dict[str, Any],
    output_path: str,
) -> None:
    """
    绘制收敛分析综合图

    Args:
        result_dict: 包含以下键:
            - r_eff_by_model: {model: r_eff}
            - distance_matrix: n×n 矩阵
            - model_names: 标签列表
        output_path: 输出路径
    """
    check_mpl_available()

    fig, axes = plt.subplots(1, 2, figsize=(16, 6))

    # 左图：有效秩对比
    r_eff_data = result_dict.get("r_eff_by_model", {})
    if r_eff_data:
        labels = list(r_eff_data.keys())
        values = [r_eff_data[l] for l in labels]

        axes[0].bar(range(len(labels)), values, color="steelblue", alpha=0.8)
        axes[0].set_xticks(range(len(labels)))
        axes[0].set_xticklabels(labels, rotation=45, ha="right")
        axes[0].set_title("Effective Rank by Model")
        axes[0].set_ylabel("Effective Rank")
        axes[0].grid(True, alpha=0.3)

    # 右图：距离热力图
    dist_matrix = result_dict.get("distance_matrix")
    model_names = result_dict.get("model_names", [])

    if dist_matrix is not None and len(model_names) > 0:
        sns.heatmap(
            dist_matrix,
            annot=True,
            fmt=".3f",
            cmap="viridis",
            xticklabels=model_names,
            yticklabels=model_names,
            ax=axes[1],
            cbar_kws={"label": "W1 Distance"},
        )
        axes[1].set_title("Cross-Model Spectrum Distance")

    plt.tight_layout()
    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
    plt.savefig(output_path, dpi=150, bbox_inches="tight")
    plt.close(fig)


def generate_all_visualizations(
    result_dir: str,
    output_dir: str,
) -> Dict[str, str]:
    """
    从实验结果生成所有可视化图

    Args:
        result_dir: 实验结果目录
        output_dir: 输出目录

    Returns:
        {plot_name: output_path, ...}
    """
    check_mpl_available()

    # 加载实验结果
    with open(Path(result_dir) / "spectrum_results.json", "r") as f:
        spectrum_results = json.load(f)

    with open(Path(result_dir) / "config.json", "r") as f:
        config = json.load(f)

    # 加载距离矩阵
    dist_dir = Path(result_dir) / "distance_matrices"
    distance_matrices = {}
    for path in dist_dir.glob("*.npy"):
        domain = path.stem.replace("_w1_matrix", "")
        distance_matrices[domain] = np.load(path)

    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)

    generated = {}

    # 1. 每个领域的距离热力图
    for domain, matrix in distance_matrices.items():
        plot_path = str(output_path / f"heatmap_{domain}.png")
        plot_heatmap(matrix, config["model_names"], title=f"{domain} - Spectrum Distance", output_path=plot_path)
        generated[f"heatmap_{domain}"] = plot_path

    # 2. 有效秩柱状图（按领域分组）
    for domain in config["domains"]:
        r_eff_data = {}
        for key, data in spectrum_results.items():
            if data["domain"] == domain:
                r_eff_data[data["model_name"]] = data["r_eff"]

        if r_eff_data:
            plot_path = str(output_path / f"bar_r_eff_{domain}.png")
            plot_bar_chart(r_eff_data, title=f"Effective Rank - {domain}", output_path=plot_path)
            generated[f"bar_r_eff_{domain}"] = plot_path

    # 3. 谱分布对比图（每个领域）
    for domain in config["domains"]:
        prob_dict = {}
        for key, data in spectrum_results.items():
            if data["domain"] == domain:
                prob_dict[data["model_name"]] = data["prob"]

        if prob_dict:
            plot_path = str(output_path / f"spectrum_{domain}.png")
            plot_spectrum_comparison(prob_dict, title=f"Spectrum Distribution - {domain}", output_path=plot_path)
            generated[f"spectrum_{domain}"] = plot_path

    # 4. 综合收敛分析图
    for domain, matrix in distance_matrices.items():
        r_eff_data = {}
        for key, data in spectrum_results.items():
            if data["domain"] == domain:
                r_eff_data[data["model_name"]] = data["r_eff"]

        result_dict = {
            "r_eff_by_model": r_eff_data,
            "distance_matrix": matrix,
            "model_names": config["model_names"],
        }

        plot_path = str(output_path / f"convergence_{domain}.png")
        plot_convergence_analysis(result_dict, plot_path)
        generated[f"convergence_{domain}"] = plot_path

    return generated


# 命令行入口
if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="生成实验可视化图")
    parser.add_argument("--result-dir", type=str, required=True, help="实验结果目录")
    parser.add_argument("--output-dir", type=str, default="experiments/output/figures", help="输出目录")

    args = parser.parse_args()
    generated = generate_all_visualizations(args.result_dir, args.output_dir)

    print(f"已生成 {len(generated)} 个可视化图:")
    for name, path in generated.items():
        print(f"  - {name}: {path}")