# paint/rose_chart.py
"""
玫瑰图（南丁格尔玫瑰图）模块
用于展示周期性和方向性数据
"""

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.patches import Rectangle, Wedge

def rose_chart(values=None, categories=None, colors=None,
               title="Rose Chart (Nightingale Rose)",
               cmap='viridis', edgecolor='white', linewidth=1,
               show_labels=True, label_fontsize=10,
               show_values=True, value_fontsize=9,
               start_angle=90, normalize=True,
               direction='clockwise', max_radius=None,
               highlight_indices=None, highlight_color='red'):
    """
    绘制玫瑰图（南丁格尔玫瑰图）

    参数:
    values: 数值列表
    categories: 类别标签列表
    colors: 颜色列表，可选
    title: 图表标题
    cmap: 颜色映射，当colors为None时使用
    edgecolor: 扇形边框颜色
    linewidth: 边框宽度
    show_labels: 是否显示类别标签
    label_fontsize: 标签字体大小
    show_values: 是否显示数值标签
    value_fontsize: 数值标签字体大小
    start_angle: 起始角度（度）
    normalize: 是否根据数值自动调整半径
    direction: 方向，'clockwise'（顺时针）或'counterclockwise'（逆时针）
    max_radius: 最大半径，如果为None则自动计算
    highlight_indices: 要高亮的扇形索引列表
    highlight_color: 高亮颜色

    返回:
    fig, ax, wedges
    """
    # 生成默认数据
    if values is None:
        values = [10, 15, 7, 12, 9, 14, 8, 11]
        categories = [f'Category {i+1}' for i in range(len(values))]

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

    if len(values) != len(categories):
        raise ValueError(f"数值数量({len(values)})与类别数量({len(categories)})不匹配")

    n = len(values)

    # 转换起始角度为弧度
    start_angle_rad = np.deg2rad(start_angle)

    # 计算每个扇形的角度
    theta = 2 * np.pi / n

    # 确定半径
    if normalize:
        max_value = max(values)
        radii = [v / max_value for v in values]
        if max_radius is None:
            max_radius = 1.0
    else:
        radii = values
        if max_radius is None:
            max_radius = max(values) * 1.1

    # 设置颜色
    if colors is None:
        colors = plt.cm.get_cmap(cmap)(np.linspace(0, 1, n))
    elif len(colors) < n:
        colors = colors * (n // len(colors) + 1)

    # 应用高亮
    if highlight_indices:
        for idx in highlight_indices:
            if idx < n:
                colors[idx] = highlight_color

    # 创建极坐标图
    fig, ax = plt.subplots(figsize=(10, 8), subplot_kw=dict(polar=True))

    # 设置方向
    if direction == 'clockwise':
        ax.set_theta_direction(-1)

    # 绘制每个扇形
    wedges = []
    for i, (radius, category, color) in enumerate(zip(radii, categories, colors)):
        # 计算当前扇形的角度
        angle = start_angle_rad + i * theta

        # 计算实际半径
        actual_radius = radius * max_radius if normalize else radius

        # 创建扇形
        wedge = Wedge((0, 0), actual_radius,
                     np.degrees(angle),
                     np.degrees(angle + theta),
                     facecolor=color,
                     edgecolor=edgecolor,
                     linewidth=linewidth,
                     alpha=0.8)

        ax.add_patch(wedge)
        wedges.append(wedge)

        # 添加数值标签（在扇形的中间位置）
        if show_values:
            # 计算标签位置
            label_angle = angle + theta / 2
            label_radius = actual_radius * 0.5

            # 转换坐标
            x = label_radius * np.cos(label_angle)
            y = label_radius * np.sin(label_angle)

            # 添加文本
            ax.text(label_angle, label_radius,
                   f'{values[i]}',
                   ha='center', va='center',
                   fontsize=value_fontsize,
                   fontweight='bold',
                   color='white' if np.mean(color[:3]) < 0.5 else 'black')

    # 添加类别标签
    if show_labels:
        for i, category in enumerate(categories):
            angle = start_angle_rad + i * theta + theta / 2
            # 标签位置稍微超出最大半径
            label_radius = max_radius * 1.15 if normalize else max(values) * 1.15

            ax.text(angle, label_radius, category,
                   ha='center', va='center',
                   fontsize=label_fontsize,
                   fontweight='bold')

    # 设置极坐标轴
    ax.set_xticks(np.linspace(0, 2*np.pi, n, endpoint=False))
    ax.set_xticklabels([])  # 隐藏默认的刻度标签

    # 设置半径刻度
    if normalize:
        # 显示百分比
        tick_values = np.linspace(0, max_radius, 5)
        tick_labels = [f'{int(t*100)}%' for t in tick_values/max_radius]
        ax.set_yticks(tick_values)
        ax.set_yticklabels(tick_labels, fontsize=9)
    else:
        # 显示实际值
        ax.set_yticks(np.linspace(0, max_radius, 5))

    # 设置标题
    plt.title(title, fontsize=16, fontweight='bold', pad=30)

    # 移除极坐标网格
    ax.grid(False)

    plt.tight_layout()
    return fig, ax, wedges


def stacked_rose_chart(data_matrix=None, categories=None, series_labels=None,
                       title="Stacked Rose Chart", cmap='tab20c',
                       start_angle=90, normalize=True, max_radius=1.0,
                       show_legend=True):
    """
    绘制堆叠玫瑰图

    参数:
    data_matrix: 数据矩阵 (n_categories, n_series)
    categories: 类别标签
    series_labels: 系列标签
    title: 图表标题
    cmap: 颜色映射
    start_angle: 起始角度
    normalize: 是否归一化
    max_radius: 最大半径
    show_legend: 是否显示图例
    """
    # 生成默认数据
    if data_matrix is None:
        np.random.seed(42)
        n_categories = 8
        n_series = 3
        data_matrix = np.random.randint(5, 20, (n_categories, n_series))
        categories = [f'Cat {i+1}' for i in range(n_categories)]
        series_labels = [f'Series {i+1}' for i in range(n_series)]

    n_categories, n_series = data_matrix.shape

    if categories is None:
        categories = [f'Category {i+1}' for i in range(n_categories)]
    if series_labels is None:
        series_labels = [f'Series {i+1}' for i in range(n_series)]

    # 转换起始角度为弧度
    start_angle_rad = np.deg2rad(start_angle)

    # 计算每个扇形的角度
    theta = 2 * np.pi / n_categories

    # 归一化数据（如果需要）
    if normalize:
        row_sums = data_matrix.sum(axis=1)
        normalized_data = data_matrix / row_sums[:, np.newaxis]
        data_to_plot = normalized_data * max_radius
    else:
        data_to_plot = data_matrix

    # 创建颜色
    colors = plt.cm.get_cmap(cmap)(np.linspace(0, 1, n_series))

    # 创建极坐标图
    fig, ax = plt.subplots(figsize=(12, 9), subplot_kw=dict(polar=True))

    # 绘制堆叠扇形
    wedges_by_series = [[] for _ in range(n_series)]

    # 对于每个类别
    for cat_idx in range(n_categories):
        current_radius = 0

        # 对于每个系列（从内到外绘制）
        for series_idx in range(n_series):
            radius_to_add = data_to_plot[cat_idx, series_idx]
            angle = start_angle_rad + cat_idx * theta

            # 创建扇形
            wedge = Wedge((0, 0), current_radius + radius_to_add,
                         np.degrees(angle), np.degrees(angle + theta),
                         facecolor=colors[series_idx],
                         edgecolor='white',
                         linewidth=1,
                         alpha=0.7)

            ax.add_patch(wedge)
            wedges_by_series[series_idx].append(wedge)

            # 添加标签（在最外层扇形上）
            if series_idx == n_series - 1:
                # 添加类别标签
                label_angle = angle + theta / 2
                label_radius = current_radius + radius_to_add + max_radius * 0.05

                ax.text(label_angle, label_radius,
                       categories[cat_idx],
                       ha='center', va='center',
                       fontsize=10, fontweight='bold')

            current_radius += radius_to_add

    # 设置极坐标轴
    ax.set_xticks(np.linspace(0, 2*np.pi, n_categories, endpoint=False))
    ax.set_xticklabels([])

    # 设置半径刻度
    if normalize:
        tick_values = np.linspace(0, max_radius, 5)
        tick_labels = [f'{int(t*100)}%' for t in tick_values/max_radius]
        ax.set_yticks(tick_values)
        ax.set_yticklabels(tick_labels, fontsize=9)
    else:
        max_value = data_to_plot.sum(axis=1).max()
        ax.set_yticks(np.linspace(0, max_value, 5))

    # 设置标题
    plt.title(title, fontsize=18, fontweight='bold', pad=40)

    # 添加图例
    if show_legend:
        from matplotlib.patches import Patch
        legend_elements = [Patch(facecolor=colors[i],
                                label=series_labels[i],
                                alpha=0.7)
                          for i in range(n_series)]

        ax.legend(handles=legend_elements,
                 loc='upper right',
                 bbox_to_anchor=(1.3, 1.0),
                 fontsize=10,
                 title='Series')

    # 移除极坐标网格
    ax.grid(False)

    plt.tight_layout()
    return fig, ax, wedges_by_series


def wind_rose(wind_direction, wind_speed, direction_bins=16,
              speed_bins=None, title="Wind Rose", cmap='coolwarm',
              show_legend=True, calm_threshold=0.5):
    """
    绘制风向玫瑰图（专门用于风向风速数据）

    参数:
    wind_direction: 风向数据（度，0-360）
    wind_speed: 风速数据
    direction_bins: 方向分箱数量（通常8, 16或32）
    speed_bins: 风速分箱边界，如[0, 5, 10, 15, 20]
    title: 图表标题
    cmap: 颜色映射
    show_legend: 是否显示图例
    calm_threshold: 静风阈值（低于此值被认为是静风）
    """
    # 确保数据是numpy数组
    wind_direction = np.array(wind_direction)
    wind_speed = np.array(wind_speed)

    if len(wind_direction) != len(wind_speed):
        raise ValueError("风向和风速数据长度必须相同")

    # 移除静风数据（可选）
    mask = wind_speed >= calm_threshold
    wind_direction = wind_direction[mask]
    wind_speed = wind_speed[mask]

    # 设置默认风速分箱
    if speed_bins is None:
        max_speed = wind_speed.max()
        speed_bins = np.linspace(0, max_speed, 6)

    n_speed_bins = len(speed_bins) - 1
    n_dir_bins = direction_bins

    # 计算方向分箱
    dir_bin_edges = np.linspace(0, 360, n_dir_bins + 1)
    dir_bin_centers = (dir_bin_edges[:-1] + dir_bin_edges[1:]) / 2

    # 计算每个方向和风速组合的频率
    frequency = np.zeros((n_dir_bins, n_speed_bins))

    for i in range(n_dir_bins):
        # 选择当前方向范围内的数据
        dir_mask = (wind_direction >= dir_bin_edges[i]) & (wind_direction < dir_bin_edges[i+1])
        if dir_mask.any():
            dir_speeds = wind_speed[dir_mask]

            # 计算风速分布
            for j in range(n_speed_bins):
                speed_mask = (dir_speeds >= speed_bins[j]) & (dir_speeds < speed_bins[j+1])
                frequency[i, j] = speed_mask.sum()

    # 转换为百分比
    frequency_percent = frequency / len(wind_direction) * 100

    # 创建颜色
    colors = plt.cm.get_cmap(cmap)(np.linspace(0, 1, n_speed_bins))

    # 创建极坐标图
    fig, ax = plt.subplots(figsize=(10, 8), subplot_kw=dict(polar=True))

    # 设置风向（0度为正北，顺时针）
    ax.set_theta_zero_location('N')
    ax.set_theta_direction(-1)  # 顺时针

    # 绘制堆叠条形
    bottom = np.zeros(n_dir_bins)

    bars_by_speed = []
    for j in range(n_speed_bins):
        # 角度转换为弧度
        theta = np.deg2rad(dir_bin_centers)
        width = np.deg2rad(360 / n_dir_bins)

        bars = ax.bar(theta, frequency_percent[:, j],
                     width=width, bottom=bottom,
                     color=colors[j], edgecolor='white',
                     linewidth=0.5, alpha=0.8,
                     label=f'{speed_bins[j]:.1f}-{speed_bins[j+1]:.1f} m/s')

        bottom += frequency_percent[:, j]
        bars_by_speed.append(bars)

    # 设置方向标签
    dir_labels = ['N', 'NNE', 'NE', 'ENE', 'E', 'ESE', 'SE', 'SSE',
                 'S', 'SSW', 'SW', 'WSW', 'W', 'WNW', 'NW', 'NNW']

    if n_dir_bins == 8:
        dir_labels = ['N', 'NE', 'E', 'SE', 'S', 'SW', 'W', 'NW']
    elif n_dir_bins == 32:
        dir_labels = ['N', 'NNE', 'NE', 'ENE', 'E', 'ESE', 'SE', 'SSE',
                     'S', 'SSW', 'SW', 'WSW', 'W', 'WNW', 'NW', 'NNW'] * 2

    ax.set_xticks(np.deg2rad(dir_bin_centers))
    ax.set_xticklabels(dir_labels[:n_dir_bins], fontsize=10)

    # 设置半径标签（百分比）
    max_percent = frequency_percent.sum(axis=1).max()
    ax.set_yticks(np.linspace(0, max_percent, 5))
    ax.set_yticklabels([f'{p:.1f}%' for p in np.linspace(0, max_percent, 5)], fontsize=9)

    # 设置标题
    plt.title(title, fontsize=16, fontweight='bold', pad=30)

    # 添加图例
    if show_legend:
        ax.legend(title='Wind Speed (m/s)',
                 loc='upper right',
                 bbox_to_anchor=(1.3, 1.0),
                 fontsize=9)

    # 添加统计信息
    calm_percent = 100 - mask.sum() / len(mask) * 100
    stats_text = f'N = {len(wind_direction)}\nMean Speed = {wind_speed.mean():.2f} m/s\nCalm (<{calm_threshold}m/s) = {calm_percent:.1f}%'

    props = dict(boxstyle='round', facecolor='wheat', alpha=0.5)
    ax.text(0.02, 0.98, stats_text, transform=ax.transAxes, fontsize=9,
            verticalalignment='top', bbox=props)

    plt.tight_layout()
    return fig, ax, bars_by_speed


def create_sample_wind_data(n=1000):
    """创建示例风向风速数据"""
    np.random.seed(42)

    # 风向（集中在西风带）
    wind_dir = np.random.vonmises(3.14, 2, n)  # 集中在180度左右（南风）
    wind_dir = np.degrees(wind_dir) % 360

    # 风速（威布尔分布）
    wind_speed = np.random.weibull(2, n) * 10

    return wind_dir, wind_speed