import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
import numpy as np

def plot_histogram(df, column, bins=30, title=None):
    """绘制直方图"""
    plt.figure(figsize=(10, 6))
    sns.histplot(df[column], bins=bins, kde=True)
    plt.title(title or f'{column} 分布')
    plt.xlabel(column)
    plt.ylabel('频率')
    plt.tight_layout()
    return plt

def plot_scatter(df, x, y, hue=None, title=None):
    """绘制散点图"""
    plt.figure(figsize=(10, 6))
    sns.scatterplot(data=df, x=x, y=y, hue=hue)
    plt.title(title or f'{x} vs {y}')
    plt.tight_layout()
    return plt

def plot_bar(df, x, y, title=None):
    """绘制柱状图"""
    plt.figure(figsize=(10, 6))
    sns.barplot(data=df, x=x, y=y)
    plt.title(title or f'{y} by {x}')
    plt.xticks(rotation=45)
    plt.tight_layout()
    return plt

def plot_box(df, x, y=None, title=None):
    """绘制箱线图"""
    plt.figure(figsize=(10, 6))
    if y:
        sns.boxplot(data=df, x=x, y=y)
    else:
        sns.boxplot(data=df[x])
    plt.title(title or f'{y or x} 分布')
    plt.xticks(rotation=45)
    plt.tight_layout()
    return plt

def plot_correlation_heatmap(df, title=None):
    """绘制相关性热图"""
    plt.figure(figsize=(12, 10))
    corr = df.corr()
    sns.heatmap(corr, annot=True, cmap='coolwarm', vmin=-1, vmax=1)
    plt.title(title or '相关性矩阵')
    plt.tight_layout()
    return plt

def plot_time_series(df, x, y, title=None):
    """绘制时间序列图"""
    plt.figure(figsize=(12, 6))
    plt.plot(df[x], df[y])
    plt.title(title or f'{y} 时间序列')
    plt.xlabel('时间')
    plt.ylabel(y)
    plt.tight_layout()
    return plt


def plot_group_comparison(df, group_col, value_col, res,
                           figsize=(8,6), palette='Set2',
                           show_points=False, point_alpha=0.3,
                           sig_type='letters',  # 'stars' 或 'letters'
                           show_stats=True,     # 是否显示中位数、均值、标准差
                           ylabel=None, title=None):
    """
    绘制组间比较箱线图，可标注显著性、样本量及统计量。

    参数：
        df : DataFrame，原始数据（已清洗）
        group_col : str，分组列名
        value_col : str，数值列名
        res : compare_groups 返回的结果字典
        figsize : 图尺寸
        palette : 颜色主题
        show_points : 是否显示原始数据点（抖动散点）
        point_alpha : 点透明度
        sig_type : 'stars' 用星号，'letters' 用字母标记
        show_stats : 是否显示每组的中位数、均值±标准差
        ylabel : y轴标签，默认使用 value_col
        title : 图标题
    
    作者：Sissi
    """
    # 复制数据并确保分组列为字符串
    plot_df = df[[group_col, value_col]].copy()
    plot_df[group_col] = plot_df[group_col].astype(str)

    # 获取组顺序（按中位数排序）
    medians = plot_df.groupby(group_col)[value_col].median().sort_values()
    group_order = medians.index.tolist()

    # 创建图形
    fig, ax = plt.subplots(figsize=figsize)

    # 绘制箱线图
    sns.boxplot(data=plot_df, x=group_col, y=value_col,
                order=group_order, palette=palette,
                width=0.5, fliersize=0, ax=ax)

    # 可选：添加抖动散点
    if show_points:
        sns.stripplot(data=plot_df, x=group_col, y=value_col,
                      order=group_order, color='black', alpha=point_alpha,
                      size=2, jitter=0.2, ax=ax)

    # 标注显著性
    if res['test']['significant'] and res['posthoc'] is not None:
        posthoc = res['posthoc']
        group_names = group_order

        if sig_type == 'stars':
            # 此处略，可自行实现星号标注
            pass
        else:
            # 字母标记法
            sig_matrix = posthoc < 0.05
            groups = group_names
            n = len(groups)
            labels = [''] * n
            current_letter = 'a'
            for i in range(n):
                if labels[i] == '':
                    labels[i] = current_letter
                    for j in range(i+1, n):
                        if not sig_matrix.loc[groups[i], groups[j]]:
                            if labels[j] == '':
                                labels[j] = current_letter
                    current_letter = chr(ord(current_letter) + 1)

            # 添加字母标记
            x_positions = [ax.get_xticks()[group_order.index(g)] for g in groups]
            y_max = plot_df[value_col].max()
            y_range = plot_df[value_col].max() - plot_df[value_col].min()
            y_offset = y_range * 0.05
            for i, label in enumerate(labels):
                ax.text(x_positions[i], y_max + y_offset, label,
                        ha='center', va='bottom', fontsize=12, fontweight='bold')

    # 添加样本量标注（底部）
    n_counts = plot_df.groupby(group_col).size()
    x_positions = [ax.get_xticks()[group_order.index(g)] for g in group_order]
    y_min = plot_df[value_col].min()
    y_range = plot_df[value_col].max() - plot_df[value_col].min()
    y_text = y_min - y_range * 0.1
    for i, g in enumerate(group_order):
        ax.text(x_positions[i], y_text, f'n={n_counts[g]}',
                ha='center', va='top', fontsize=9, color='gray')

    # 添加统计量标注（中位数、均值±标准差）
    if show_stats:
        # 从 res['descriptive'] 提取统计量
        desc = res['descriptive'].loc[group_order]  # 按顺序排列
        # 中位数标注：放在箱子顶部上方一点
        y_max_plot = ax.get_ylim()[1]
        y_range_plot = y_max_plot - ax.get_ylim()[0]
        y_median_pos = y_max_plot + y_range_plot * 0.05  # 在显著性字母下方
        for i, g in enumerate(group_order):
            median_val = desc.loc[g, 'median']
            ax.text(x_positions[i], y_median_pos, f'med={median_val:.1f}',
                    ha='center', va='bottom', fontsize=8, color='darkblue')

        # 均值±标准差标注：放在箱子底部下方，样本量上方
        y_mean_pos = y_min - y_range * 0.2  # 在样本量下方再留空间
        for i, g in enumerate(group_order):
            mean_val = desc.loc[g, 'mean']
            std_val = desc.loc[g, 'std']
            ax.text(x_positions[i], y_mean_pos, r'$\mu$={mean_val:.1f}\n$\sigma$={std_val:.1f}'.format(mean_val=mean_val, std_val=std_val),
                    ha='center', va='top', fontsize=7, color='darkgreen', linespacing=0.8)
        # 调整 y 轴范围以容纳更多底部标注
        ax.set_ylim(y_mean_pos - y_range*0.1, y_max_plot + y_range_plot*0.15)
    else:
        # 不显示统计量时，仅调整 y 轴范围以容纳样本量和显著性
        ylim = ax.get_ylim()
        ax.set_ylim(y_text - y_range*0.05, ylim[1] + y_range*0.15)

    # 设置标签
    ax.set_xlabel(group_col, fontsize=12)
    ax.set_ylabel(ylabel if ylabel else value_col, fontsize=12)
    if title:
        ax.set_title(title, fontsize=14)

    # 添加网格线
    ax.grid(axis='y', linestyle='--', alpha=0.5)

    plt.tight_layout()
    return fig, ax
