"""
Benchmark 结果可视化模块
"""

import os
from typing import List, Dict, Any
from datetime import datetime

import numpy as np

from .results import TestResult
from .utils import ensure_dir, setup_logging

logger = setup_logging()


def generate_visualizations(results: List[TestResult], output_dir: str) -> List[str]:
    """生成所有可视化图表"""
    try:
        import matplotlib.pyplot as plt
        import matplotlib
        matplotlib.use('Agg')  # 非交互式后端
        plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
        plt.rcParams['axes.unicode_minus'] = False
    except ImportError:
        logger.warning("matplotlib 未安装，跳过可视化")
        return []
    
    ensure_dir(output_dir)
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    generated_files = []
    
    # 1. GPU 利用率 vs 摄像头数量
    fig_path = _plot_gpu_vs_cameras(results, output_dir, timestamp)
    if fig_path:
        generated_files.append(fig_path)
    
    # 2. 吞吐量 vs Batch Size
    fig_path = _plot_throughput_vs_batch(results, output_dir, timestamp)
    if fig_path:
        generated_files.append(fig_path)
    
    # 3. 延迟分布
    fig_path = _plot_latency_distribution(results, output_dir, timestamp)
    if fig_path:
        generated_files.append(fig_path)
    
    # 4. 实时性热力图
    fig_path = _plot_realtime_heatmap(results, output_dir, timestamp)
    if fig_path:
        generated_files.append(fig_path)
    
    # 5. GPU 饱和点分析
    fig_path = _plot_saturation_analysis(results, output_dir, timestamp)
    if fig_path:
        generated_files.append(fig_path)
    
    logger.info(f"生成 {len(generated_files)} 个可视化图表")
    return generated_files


def _plot_gpu_vs_cameras(results: List[TestResult], output_dir: str, timestamp: str) -> str:
    """GPU 利用率 vs 摄像头数量"""
    import matplotlib.pyplot as plt
    
    fig, ax = plt.subplots(figsize=(10, 6))
    
    # 按 batch size 分组
    batch_sizes = sorted(set(r.batch_size for r in results))
    camera_counts = sorted(set(r.num_cameras for r in results))
    
    for bs in batch_sizes:
        bs_results = [r for r in results if r.batch_size == bs]
        cameras = []
        gpu_utils = []
        for cam in camera_counts:
            cam_results = [r for r in bs_results if r.num_cameras == cam]
            if cam_results:
                cameras.append(cam)
                gpu_utils.append(np.mean([r.gpu_utilization_avg for r in cam_results]))
        
        if cameras:
            ax.plot(cameras, gpu_utils, marker='o', label=f'Batch={bs}')
    
    ax.axhline(y=85, color='r', linestyle='--', label='饱和阈值 (85%)')
    ax.set_xlabel('摄像头数量')
    ax.set_ylabel('GPU 利用率 (%)')
    ax.set_title('GPU 利用率 vs 摄像头数量')
    ax.legend()
    ax.grid(True, alpha=0.3)
    ax.set_ylim(0, 100)
    
    fig_path = os.path.join(output_dir, f'gpu_vs_cameras_{timestamp}.png')
    plt.savefig(fig_path, dpi=150, bbox_inches='tight')
    plt.close()
    
    return fig_path


def _plot_throughput_vs_batch(results: List[TestResult], output_dir: str, timestamp: str) -> str:
    """吞吐量 vs Batch Size"""
    import matplotlib.pyplot as plt
    
    fig, ax = plt.subplots(figsize=(10, 6))
    
    batch_sizes = sorted(set(r.batch_size for r in results))
    camera_counts = sorted(set(r.num_cameras for r in results))
    
    for cam in camera_counts:
        cam_results = [r for r in results if r.num_cameras == cam]
        batches = []
        throughputs = []
        for bs in batch_sizes:
            bs_results = [r for r in cam_results if r.batch_size == bs]
            if bs_results:
                batches.append(bs)
                throughputs.append(np.mean([r.total_throughput_fps for r in bs_results]))
        
        if batches:
            ax.plot(batches, throughputs, marker='s', label=f'{cam} 路摄像头')
    
    ax.set_xlabel('Batch Size')
    ax.set_ylabel('总吞吐量 (FPS)')
    ax.set_title('吞吐量 vs Batch Size')
    ax.legend()
    ax.grid(True, alpha=0.3)
    ax.set_xticks(batch_sizes)
    
    fig_path = os.path.join(output_dir, f'throughput_vs_batch_{timestamp}.png')
    plt.savefig(fig_path, dpi=150, bbox_inches='tight')
    plt.close()
    
    return fig_path


def _plot_latency_distribution(results: List[TestResult], output_dir: str, timestamp: str) -> str:
    """延迟分布对比"""
    import matplotlib.pyplot as plt
    
    fig, axes = plt.subplots(1, 2, figsize=(14, 5))
    
    batch_sizes = sorted(set(r.batch_size for r in results))
    
    # 左图: 平均延迟
    ax1 = axes[0]
    x = np.arange(len(batch_sizes))
    width = 0.35
    
    avg_latencies = [np.mean([r.avg_latency_ms for r in results if r.batch_size == bs]) for bs in batch_sizes]
    p95_latencies = [np.mean([r.p95_latency_ms for r in results if r.batch_size == bs]) for bs in batch_sizes]
    
    ax1.bar(x - width/2, avg_latencies, width, label='平均延迟')
    ax1.bar(x + width/2, p95_latencies, width, label='P95 延迟')
    ax1.set_xlabel('Batch Size')
    ax1.set_ylabel('延迟 (ms)')
    ax1.set_title('延迟 vs Batch Size')
    ax1.set_xticks(x)
    ax1.set_xticklabels(batch_sizes)
    ax1.legend()
    ax1.grid(True, alpha=0.3, axis='y')
    
    # 右图: 按摄像头数量
    ax2 = axes[1]
    camera_counts = sorted(set(r.num_cameras for r in results))
    
    avg_by_cam = [np.mean([r.avg_latency_ms for r in results if r.num_cameras == cam]) for cam in camera_counts]
    p95_by_cam = [np.mean([r.p95_latency_ms for r in results if r.num_cameras == cam]) for cam in camera_counts]
    
    x2 = np.arange(len(camera_counts))
    ax2.bar(x2 - width/2, avg_by_cam, width, label='平均延迟')
    ax2.bar(x2 + width/2, p95_by_cam, width, label='P95 延迟')
    ax2.set_xlabel('摄像头数量')
    ax2.set_ylabel('延迟 (ms)')
    ax2.set_title('延迟 vs 摄像头数量')
    ax2.set_xticks(x2)
    ax2.set_xticklabels(camera_counts)
    ax2.legend()
    ax2.grid(True, alpha=0.3, axis='y')
    
    plt.tight_layout()
    fig_path = os.path.join(output_dir, f'latency_distribution_{timestamp}.png')
    plt.savefig(fig_path, dpi=150, bbox_inches='tight')
    plt.close()
    
    return fig_path


def _plot_realtime_heatmap(results: List[TestResult], output_dir: str, timestamp: str) -> str:
    """实时性热力图"""
    import matplotlib.pyplot as plt
    
    batch_sizes = sorted(set(r.batch_size for r in results))
    camera_counts = sorted(set(r.num_cameras for r in results))
    
    # 创建矩阵
    matrix = np.zeros((len(camera_counts), len(batch_sizes)))
    
    for i, cam in enumerate(camera_counts):
        for j, bs in enumerate(batch_sizes):
            matching = [r for r in results if r.num_cameras == cam and r.batch_size == bs]
            if matching:
                # 实时性比例
                realtime_ratio = sum(1 for r in matching if r.is_realtime_capable) / len(matching)
                matrix[i, j] = realtime_ratio * 100
    
    fig, ax = plt.subplots(figsize=(10, 8))
    
    im = ax.imshow(matrix, cmap='RdYlGn', aspect='auto', vmin=0, vmax=100)
    
    ax.set_xticks(np.arange(len(batch_sizes)))
    ax.set_yticks(np.arange(len(camera_counts)))
    ax.set_xticklabels(batch_sizes)
    ax.set_yticklabels(camera_counts)
    ax.set_xlabel('Batch Size')
    ax.set_ylabel('摄像头数量')
    ax.set_title('实时性热力图 (绿色=可实时, 红色=不可实时)')
    
    # 添加数值标注
    for i in range(len(camera_counts)):
        for j in range(len(batch_sizes)):
            text = ax.text(j, i, f'{matrix[i, j]:.0f}%',
                          ha='center', va='center', color='black', fontsize=9)
    
    plt.colorbar(im, ax=ax, label='实时性比例 (%)')
    
    fig_path = os.path.join(output_dir, f'realtime_heatmap_{timestamp}.png')
    plt.savefig(fig_path, dpi=150, bbox_inches='tight')
    plt.close()
    
    return fig_path


def _plot_saturation_analysis(results: List[TestResult], output_dir: str, timestamp: str) -> str:
    """GPU 饱和点分析"""
    import matplotlib.pyplot as plt
    
    fig, axes = plt.subplots(2, 2, figsize=(14, 10))
    
    camera_counts = sorted(set(r.num_cameras for r in results))
    batch_sizes = sorted(set(r.batch_size for r in results))
    
    # 1. GPU 利用率 vs 吞吐量
    ax1 = axes[0, 0]
    for bs in batch_sizes:
        bs_results = [r for r in results if r.batch_size == bs]
        throughputs = [r.total_throughput_fps for r in bs_results]
        gpu_utils = [r.gpu_utilization_avg for r in bs_results]
        ax1.scatter(throughputs, gpu_utils, label=f'Batch={bs}', alpha=0.7)
    
    ax1.axhline(y=85, color='r', linestyle='--', alpha=0.5)
    ax1.set_xlabel('吞吐量 (FPS)')
    ax1.set_ylabel('GPU 利用率 (%)')
    ax1.set_title('GPU 利用率 vs 吞吐量')
    ax1.legend()
    ax1.grid(True, alpha=0.3)
    
    # 2. 显存使用
    ax2 = axes[0, 1]
    for bs in batch_sizes:
        bs_results = [r for r in results if r.batch_size == bs]
        cameras = []
        mem_used = []
        for cam in camera_counts:
            cam_results = [r for r in bs_results if r.num_cameras == cam]
            if cam_results:
                cameras.append(cam)
                mem_used.append(np.mean([r.memory_used_mb for r in cam_results]))
        if cameras:
            ax2.plot(cameras, mem_used, marker='o', label=f'Batch={bs}')
    
    ax2.set_xlabel('摄像头数量')
    ax2.set_ylabel('显存使用 (MB)')
    ax2.set_title('显存使用 vs 摄像头数量')
    ax2.legend()
    ax2.grid(True, alpha=0.3)
    
    # 3. 丢帧率
    ax3 = axes[1, 0]
    for bs in batch_sizes:
        bs_results = [r for r in results if r.batch_size == bs]
        cameras = []
        drop_rates = []
        for cam in camera_counts:
            cam_results = [r for r in bs_results if r.num_cameras == cam]
            if cam_results:
                cameras.append(cam)
                drop_rates.append(np.mean([r.frame_drop_rate for r in cam_results]))
        if cameras:
            ax3.plot(cameras, drop_rates, marker='s', label=f'Batch={bs}')
    
    ax3.axhline(y=5, color='r', linestyle='--', label='阈值 (5%)')
    ax3.set_xlabel('摄像头数量')
    ax3.set_ylabel('丢帧率 (%)')
    ax3.set_title('丢帧率 vs 摄像头数量')
    ax3.legend()
    ax3.grid(True, alpha=0.3)
    
    # 4. 饱和状态统计
    ax4 = axes[1, 1]
    saturated_count = sum(1 for r in results if r.is_gpu_saturated)
    not_saturated_count = len(results) - saturated_count
    
    ax4.pie([not_saturated_count, saturated_count], 
            labels=['未饱和', '已饱和'],
            colors=['#2ecc71', '#e74c3c'],
            autopct='%1.1f%%',
            startangle=90)
    ax4.set_title(f'GPU 饱和状态分布\n(共 {len(results)} 个测试)')
    
    plt.tight_layout()
    fig_path = os.path.join(output_dir, f'saturation_analysis_{timestamp}.png')
    plt.savefig(fig_path, dpi=150, bbox_inches='tight')
    plt.close()
    
    return fig_path