import pandas as pd
import numpy as np
from scipy.stats import mannwhitneyu, levene
from statsmodels.stats.weightstats import CompareMeans, DescrStatsW
from statsmodels.formula.api import mixedlm
import os

# 读取CSV文件
category_means = pd.read_csv('整理后的数据_填补后.csv')


# 定义函数将数据分为A组和B组
def group_by_name_start(df):
    df['分组'] = df['姓名'].apply(lambda x: 'A组' if x.startswith('A') else ('B组' if x.startswith('B') else '其他'))
    return df[df['分组'] != '其他']  # 仅保留A组和B组的数据


# 分组
grouped_df = group_by_name_start(category_means)

# 获取所有特征列（排除‘姓名’和‘分组’列）
feature_columns = [col for col in grouped_df.columns if col not in ['姓名', '分组', 'Category']]

# 分别提取A组和B组数据
a_group = grouped_df[grouped_df['分组'] == 'A组']
b_group = grouped_df[grouped_df['分组'] == 'B组']

# 创建输出目录
output_dir = '分析结果'
if not os.path.exists(output_dir):
    os.makedirs(output_dir)

# 存储中位数和方差的比较结果
median_results = []
variance_results = []

# 存储Cohen's d和Z检验的结果
effect_size_results = []

# 存储混合效应模型的结果
mixed_model_results = []

# 进行中位数和方差比较、亚组分析和混合效应模型分析
for feature in feature_columns:
    print("================================================")
    # 中位数比较（Mann-Whitney U 检验）
    stat_median, p_median = mannwhitneyu(a_group[feature].dropna(), b_group[feature].dropna())
    median_results.append((feature, stat_median, p_median))

    # 方差比较（Levene 检验）
    stat_var, p_var = levene(a_group[feature].dropna(), b_group[feature].dropna(), center='median')
    variance_results.append((feature, stat_var, p_var))

    print(f'{feature} 中位数比较: Mann-Whitney U 统计量 = {stat_median}, p值 = {p_median}')
    print(f'{feature} 方差比较: Levene 统计量 = {stat_var}, p值 = {p_var}\n')

    # 计算Cohen's d
    cohen_d = (a_group[feature].mean() - b_group[feature].mean()) / np.sqrt(
        (a_group[feature].var() + b_group[feature].var()) / 2)

    # 进行Z检验
    desc_a = DescrStatsW(a_group[feature].dropna())
    desc_b = DescrStatsW(b_group[feature].dropna())
    cm = CompareMeans(desc_a, desc_b)
    z_stat, p_z = cm.ztest_ind(usevar='pooled')

    effect_size_results.append((feature, cohen_d, z_stat, p_z))

    print(f'{feature} Cohen\'s d = {cohen_d}, Z检验统计量 = {z_stat}, p值 = {p_z}\n')

    # 构建混合效应模型，确保公式中的列名正确处理
    try:
        model = mixedlm(f'Q("{feature}") ~ 分组', grouped_df, groups=grouped_df['姓名'])
        mixed_model_fit = model.fit(reml=False, method='nm')

        # 获取固定效应系数和随机效应方差
        fixed_effect = mixed_model_fit.params.get('分组[T.B组]', np.nan)
        random_effect_variance = mixed_model_fit.cov_re.iloc[0, 0] if not mixed_model_fit.cov_re.empty else np.nan

        mixed_model_results.append((feature, fixed_effect, random_effect_variance))

        print(f'{feature} 固定效应系数 (β_1) = {fixed_effect}')
        print(f'{feature} 随机效应方差 = {random_effect_variance}\n')

    except Exception as e:
        print(f"特征 {feature} 的混合效应模型无法拟合。错误: {str(e)}")
        mixed_model_results.append((feature, np.nan, np.nan))

    print("================================================")

# 将中位数和方差的比较结果保存到CSV
pd.DataFrame(median_results, columns=['特征', 'Mann-Whitney U 统计量', 'p值']).to_csv(
    os.path.join(output_dir, '中位数比较结果.csv'), index=False)
pd.DataFrame(variance_results, columns=['特征', 'Levene 统计量', 'p值']).to_csv(
    os.path.join(output_dir, '方差比较结果.csv'), index=False)

# 将亚组分析的效果量和Z检验结果保存到CSV
pd.DataFrame(effect_size_results, columns=['特征', 'Cohen\'s d', 'Z检验统计量', 'p值']).to_csv(
    os.path.join(output_dir, '效果量分析结果.csv'), index=False)

# 将混合效应模型的结果保存到CSV
pd.DataFrame(mixed_model_results, columns=['特征', '固定效应系数 (β_1)', '随机效应方差']).to_csv(
    os.path.join(output_dir, '混合效应模型结果.csv'), index=False)