import pandas as pd
import numpy as np
import pyDOE2
import statsmodels.api as sm
from statsmodels.formula.api import ols
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm

# 三个因子和响应值的具体的因子名称
factor1_name  = 'Glucose'
factor2_name  = 'Peptone'
factor3_name  = 'Time'
response_name = 'ergosterol'

# 三个因子的范围，一定要把高、中、低都写上
factor1_levels = [50, 65, 80]
factor2_levels = [20, 30, 40]
factor3_levels = [9,  11, 13]

# 实验结果和响应值
response_result = [3.82, 4.05, 3.98, 4.28, 3.75, 4.02, 3.88, 4.12, 3.92, 4.15, 4.08, 4.25, 4.52, 4.48, 4.51, 4.49, 4.5]


# --- 第1部分：生成 BBD 实验矩阵 ---
def generate_bbd(f1name, f2name, f3name, rsname, f1lelist, f2lelist, f3lelist):
    # 因素名称
    factors = [f1name, f2name, f3name]
    # 生成 BBD 矩阵 (k=3 表示3个因素)
    # BBD 会自动生成中心点重复实验，通常为 12+3=15 次或 12+5=17 次
    bbd_matrix = pyDOE2.bbdesign(3, center=5) 
    
    df = pd.DataFrame(bbd_matrix, columns=factors)
    
    print('--- 第1部分: BBD 01 矩阵 (k=5) ---')
    print(df, '\n\n\n')
    # 映射实际物理值
    level_map = {
        f1name:  f1lelist,      # -1, 0, 1
        f2name:  f2lelist,      # -1, 0, 1
        f3name:  f3lelist    # -1, 0, 1
    }
    
    real_df = df.copy()
    for f in factors:
        # 线性映射: val = center + coded * step
        center = level_map[f][1]
        step = level_map[f][2] - center
        real_df[f] = center + real_df[f] * step
        
    real_df[rsname] = np.nan # 预留响应值列
    return real_df



# --- 第2部分：模型拟合与分析 ---
def analyze_bbd(bbd_df, f1name, f2name, f3name, rsname, response_result):

    bbd_df[rsname] = response_result
    print('--- 第2部分: BBD 实验矩阵 (k=5) 含结果 ---')
    print(bbd_df)
    # 构建二次多项式公式: Y ~ X1 + X2 + X3 + X1*X2 + X1*X3 + X2*X3 + I(X1^2) + I(X2^2) + I(X3^2)
    formula = rsname + ' ~ ' + f1name + ' * ' + f2name + ' + ' + f1name + ' * ' + f3name + ' + ' + f2name + ' * ' + f3name + ' + ' + 'I(' + f1name + '**2) + I(' + f2name + '**2) + I(' + f3name + '**2)'

    model = ols(formula, data=bbd_df).fit()
    print("\n--- 第2部分: 回归模型分析报告 (ANOVA) ---")
    print(model.summary())
    print('\n------ ANOVA END ------\n\n\n')
    return model



# --- 第3部分：三维响应面可视化 ---
# 3. 绘图函数定义
def plot_rsm_3d(model, df, factor_x, factor_y, fixed_factor, fixed_val):
    """
    factor_x: X轴因子名称
    factor_y: Y轴因子名称
    fixed_factor: 被固定的因子名称
    fixed_val: 固定因子的取值（通常选中心点）
    """
 
    # 创建绘图网格
    x_range = np.linspace(df[factor_x].min(), df[factor_x].max(), 30)
    y_range = np.linspace(df[factor_y].min(), df[factor_y].max(), 30)
    X, Y = np.meshgrid(x_range, y_range)
    
    # 构建预测用的 DataFrame
    pred_df = pd.DataFrame({
        factor_x: X.ravel(),
        factor_y: Y.ravel(),
        fixed_factor: [fixed_val] * len(X.ravel())
    })
    
    # 预测 Z 轴（响应值）
    Z = model.predict(pred_df).values.reshape(X.shape)
    
    # 开始绘图
    fig = plt.figure(figsize=(10, 7))
    ax = fig.add_subplot(111, projection='3d')
    
    # 绘制曲面
    surf = ax.plot_surface(X, Y, Z, cmap=cm.viridis, edgecolor='none', alpha=0.8, antialiased=True)
    
    # 添加等高线 (投影到底面)
    ax.contour(X, Y, Z, zdir='z', offset=ax.get_zlim()[0], cmap=cm.viridis)
    
    # 设置标签
    ax.set_xlabel(f'{factor_x} (g/L)')
    ax.set_ylabel(f'{factor_y} (g/L or d)')
    ax.set_zlabel('Ergosterol (mg/g)')
    plt.title(f'Response Surface of {factor_x} and {factor_y}\n(Fixed {fixed_factor} = {fixed_val})')
    
    # 添加颜色条
    fig.colorbar(surf, shrink=0.5, aspect=10)
    
    # 调整视角
    ax.view_init(elev=25, azim=135)
    
    plt.show()




# 执行流程
bbd_dataframe = generate_bbd(factor1_name, factor2_name, factor3_name, response_name, factor1_levels, factor2_levels, factor3_levels)

fitted_model = analyze_bbd(bbd_dataframe, factor1_name, factor2_name, factor3_name, response_name, response_result)

plot_rsm_3d(fitted_model, bbd_dataframe, factor1_name, factor2_name, factor3_name, factor3_levels[1])
plot_rsm_3d(fitted_model, bbd_dataframe, factor1_name, factor3_name, factor2_name, factor2_levels[1])
plot_rsm_3d(fitted_model, bbd_dataframe, factor2_name, factor3_name, factor1_name, factor1_levels[1])