# vizpaint/rose_chart_enhanced.py
"""
增强版玫瑰图模块
提供动画、交互和统计功能的玫瑰图
"""

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.patches import Wedge
from matplotlib.animation import FuncAnimation
import matplotlib.colors as mcolors


def rose_chart_enhanced(values, categories=None, colors=None,
                        title="Enhanced Rose Chart",
                        animation=False, duration=3,
                        show_stats=True, sort_by_value=True,
                        highlight_top=3, explode=False):
    """
    增强版玫瑰图，支持动画和统计功能

    参数:
        values: 数值列表
        categories: 类别标签列表
        colors: 颜色列表，可选
        title: 图表标题
        animation: 是否启用生长动画
        duration: 动画时长（秒）
        show_stats: 是否显示统计信息
        sort_by_value: 是否按值排序
        highlight_top: 高亮前N个最大的扇形
        explode: 是否分离突出显示

    返回:
        fig, ax, wedges: 图形、坐标轴和扇形对象
        如果启用动画，额外返回 anim 对象
    """
    # 数据预处理
    values = np.array(values)
    n = len(values)

    if categories is None:
        categories = [f'Cat {i + 1}' for i in range(n)]

    # 排序（可选）
    if sort_by_value:
        sort_idx = np.argsort(values)[::-1]  # 降序
        values = values[sort_idx]
        categories = [categories[i] for i in sort_idx]

    # 颜色设置
    if colors is None:
        colors = plt.cm.Set3(np.linspace(0, 1, n))

    # 高亮设置
    if highlight_top > 0 and highlight_top < n:
        highlight_colors = list(colors)
        for i in range(min(highlight_top, n)):
            # 将前几个颜色加深
            rgb = mcolors.to_rgb(highlight_colors[i])
            highlight_colors[i] = tuple(min(c * 0.7, 1) for c in rgb)
        colors = highlight_colors

    # 爆炸效果
    if explode:
        explode_values = [0.05 if i < highlight_top else 0 for i in range(n)]
    else:
        explode_values = [0] * n

    # 创建图形
    fig, ax = plt.subplots(figsize=(12, 10), subplot_kw=dict(polar=True))

    # 计算角度和半径
    theta = np.linspace(0, 2 * np.pi, n, endpoint=False)
    width = 2 * np.pi / n
    radii = values / values.max() * 0.8  # 归一化半径

    # 绘制扇形
    wedges = []
    for i in range(n):
        wedge = Wedge((0, 0), radii[i],
                      np.degrees(theta[i]),
                      np.degrees(theta[i] + width),
                      facecolor=colors[i],
                      edgecolor='white',
                      linewidth=1.5,
                      alpha=0.85)
        ax.add_patch(wedge)
        wedges.append(wedge)

        # 添加标签
        label_angle = theta[i] + width / 2
        label_radius = radii[i] * 0.6
        ax.text(label_angle, label_radius, categories[i],
                ha='center', va='center',
                fontsize=10, fontweight='bold',
                rotation=np.degrees(label_angle) - 90)

        # 添加数值标签
        value_radius = radii[i] * 0.85
        ax.text(label_angle, value_radius, f'{values[i]:.1f}',
                ha='center', va='center',
                fontsize=9, color='white',
                fontweight='bold')

    # 添加标题
    ax.set_title(title, fontsize=18, fontweight='bold', pad=30)

    # 隐藏极坐标网格
    ax.grid(False)
    ax.set_xticklabels([])
    ax.set_yticklabels([])

    # 添加统计信息（可选）
    if show_stats:
        stats_text = (f'总数: {values.sum():.1f}\n'
                      f'均值: {values.mean():.2f}\n'
                      f'最大: {values.max():.1f}\n'
                      f'最小: {values.min():.1f}')

        # 在左上角添加文本框
        props = dict(boxstyle='round', facecolor='wheat', alpha=0.8)
        ax.text(0.02, 0.98, stats_text, transform=ax.transAxes,
                fontsize=10, verticalalignment='top',
                bbox=props)

    plt.tight_layout()

    # 动画效果（可选）
    if animation:
        def init():
            for wedge in wedges:
                wedge.set_radius(0)
            return wedges

        def update(frame):
            progress = frame / 100
            for i, wedge in enumerate(wedges):
                target_radius = radii[i]
                current_radius = target_radius * progress
                wedge.set_radius(current_radius)
            return wedges

        frames = int(duration * 30)  # 30 fps
        anim = FuncAnimation(fig, update, frames=frames,
                             init_func=init, blit=True,
                             interval=duration * 1000 / frames)

        return fig, ax, wedges, anim

    return fig, ax, wedges