import polars as pl
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LinearRegression, Ridge, Lasso
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import r2_score, mean_squared_error, mean_absolute_error
from sklearn.pipeline import Pipeline
import matplotlib.pyplot as plt
import pickle  # 添加pickle用于保存模型

# 设置matplotlib支持中文显示
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'DejaVu Sans']  # 指定默认字体
plt.rcParams['axes.unicode_minus'] = False  # 解决保存图像时负号'-'显示为方块的问题


def 加载工况整理后的数据():
    df1 = pl.read_csv("纯凝.csv")
    df2 = pl.read_csv("抽凝.csv")
    """
        df1和df2中的列名相同，均包含self.sample_intervals = {'汽轮发电机功率': 2, '汽机抽汽供热量': 20,
                                         '环境温度': 1, '凝结水补水流量': 1,
                                         '循泵运行方式数字编码': 0, '风机运行台数': 0,
                                         '冷却塔进水蝶阀开启数量': 0,
                                         '大气湿度': 5} 和背压列
        我想拟合一条背压随'汽轮发电机功率'、'汽机抽汽供热量'、'环境温度'、'大气湿度'、'循泵运行方式数字编码'、'风机运行台数'和'冷却塔进水蝶阀开启数量
        的关系曲线，应该怎么实现
    """
    # 合并两个数据框

    return df1


def 检查重复参数不同蝶阀数量():
    """
    检查df1中是否存在前4个参数相近，后2个参数完全相同但冷却塔进水蝶阀开启数量不同的工况
    前4个参数：汽轮发电机功率、汽机抽汽供热量、环境温度、大气湿度（允许相近）
    后2个参数：循泵运行方式数字编码、风机运行台数（要求完全相同）
    """
    df1 = pl.read_csv("纯凝.csv")

    # 定义参数列
    continuous_params = [  # 连续参数，允许相近
        '汽轮发电机功率',
        '汽机抽汽供热量',
        '环境温度',
        '大气湿度'
    ]

    discrete_params = [  # 离散参数，要求完全相同
        '循泵运行方式数字编码',
        '风机运行台数'
    ]

    valve_param = '冷却塔进水蝶阀开启数量'

    # 检查列是否存在
    all_params = continuous_params + discrete_params + [valve_param]
    missing_features = [col for col in all_params if col not in df1.columns]
    if missing_features:
        print(f"以下列在纯凝.csv中不存在: {missing_features}")
        return None

    # 为连续参数创建分组标签（将相近值归为一组）
    # 这里使用简单的四舍五入方法，可以根据需要调整精度
    df_with_groups = df1.with_columns([
        # 汽轮发电机功率四舍五入到最接近的0.5
        (pl.col('汽轮发电机功率') * 2).round().alias('汽轮发电机功率_分组'),
        # 汽机抽汽供热量四舍五入到整数
        pl.col('汽机抽汽供热量').round().alias('汽机抽汽供热量_分组'),
        # 环境温度四舍五入到整数
        pl.col('环境温度').round().alias('环境温度_分组'),
        # 大气湿度四舍五入到整数
        pl.col('大气湿度').round().alias('大气湿度_分组')
    ])

    # 分组列
    group_columns = [
                        '汽轮发电机功率_分组',
                        '汽机抽汽供热量_分组',
                        '环境温度_分组',
                        '大气湿度_分组'
                    ] + discrete_params

    # 按分组列分组，查看每组中'冷却塔进水蝶阀开启数量'的值
    grouped = df_with_groups.group_by(group_columns).agg([
        pl.col(valve_param).unique().sort().alias('不同蝶阀数量'),
        pl.col(valve_param).count().alias('记录数'),
        # 保留原始参数的范围以供参考
        pl.col('汽轮发电机功率').unique().sort().alias('汽轮发电机功率范围'),
        pl.col('汽机抽汽供热量').unique().sort().alias('汽机抽汽供热量范围'),
        pl.col('环境温度').unique().sort().alias('环境温度范围'),
        pl.col('大气湿度').unique().sort().alias('大气湿度范围'),
        # 保留行索引
        pl.first('汽轮发电机功率').alias('汽轮发电机功率_示例'),
        pl.first('汽机抽汽供热量').alias('汽机抽汽供热量_示例'),
        pl.first('环境温度').alias('环境温度_示例'),
        pl.first('大气湿度').alias('大气湿度_示例'),
        pl.first('循泵运行方式数字编码').alias('循泵运行方式数字编码_示例'),
        pl.first('风机运行台数').alias('风机运行台数_示例')
    ])

    # 找出有多个不同蝶阀数量的组合
    different_valve_counts = grouped.filter(
        pl.col('不同蝶阀数量').list.len() > 1
    )

    if len(different_valve_counts) > 0:
        print(f"发现 {len(different_valve_counts)} 种参数组合具有不同的冷却塔进水蝶阀开启数量:")
        for row in different_valve_counts.iter_rows(named=True):
            print(f"相近参数:")
            print(f"  汽轮发电机功率范围: {row['汽轮发电机功率范围']}")
            print(f"  汽机抽汽供热量范围: {row['汽机抽汽供热量范围']}")
            print(f"  环境温度范围: {row['环境温度范围']}")
            print(f"  大气湿度范围: {row['大气湿度范围']}")
            print(f"相同参数:")
            print(f"  循泵运行方式数字编码: {row['循泵运行方式数字编码_示例']}")
            print(f"  风机运行台数: {row['风机运行台数_示例']}")
            print(f"  不同蝶阀数量: {row['不同蝶阀数量']}")
            print(f"  记录数: {row['记录数']}")

            # 查找原始数据行
            conditions = [
                (pl.col('汽轮发电机功率_分组') == row['汽轮发电机功率_分组']),
                (pl.col('汽机抽汽供热量_分组') == row['汽机抽汽供热量_分组']),
                (pl.col('环境温度_分组') == row['环境温度_分组']),
                (pl.col('大气湿度_分组') == row['大气湿度_分组']),
                (pl.col('循泵运行方式数字编码') == row['循泵运行方式数字编码_示例']),
                (pl.col('风机运行台数') == row['风机运行台数_示例'])
            ]

            # 为每个条件添加
            filter_expr = conditions[0]
            for condition in conditions[1:]:
                filter_expr = filter_expr & condition

            matching_rows = df_with_groups.filter(filter_expr)

            print(f"  原始数据行:")
            for i, record in enumerate(matching_rows.iter_rows(named=True)):
                print(f"    记录{i + 1}:")
                # 输出所有列的值
                for col_name in df1.columns:
                    print(f"      {col_name}: {record[col_name]}")
                print()
    else:
        print("未发现前4个参数相近、后2个参数相同但冷却塔进水蝶阀开启数量不同的工况")

    return different_valve_counts


def 数据预处理(df):
    """
    数据预处理：选择特征列并处理缺失值
    """
    # 定义特征列
    feature_columns = [
        '汽轮发电机功率',
        '汽机抽汽供热量',
        '环境温度',
        '大气湿度',
        '循泵运行方式数字编码',
        '风机运行台数',
        '冷却塔进水蝶阀开启数量'
    ]

    # 选择特征列和目标列（假设背压列名为'背压'）
    target_column = '背压'  # 请根据实际数据列名调整

    # 检查列是否存在
    missing_features = [col for col in feature_columns if col not in df.columns]
    if missing_features:
        raise ValueError(f"以下列在数据中不存在: {missing_features}")

    if target_column not in df.columns:
        raise ValueError(f"目标列 '{target_column}' 在数据中不存在")

    # 选择数据
    df_clean = df.select([target_column] + feature_columns).drop_nulls()

    # 转换为pandas以使用sklearn
    X = df_clean.select(feature_columns).to_pandas()
    y = df_clean.select(target_column).to_pandas().values.ravel()

    return X, y, feature_columns


def 多种模型拟合(X, y, feature_columns):
    """
    使用多种模型进行拟合比较
    """
    # 分割数据
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.01, random_state=42)

    # 定义模型
    models = {
        'Linear Regression': LinearRegression(),
        'Ridge Regression': Ridge(alpha=1.0),
        'Lasso Regression': Lasso(alpha=0.1),
        'Random Forest': RandomForestRegressor(n_estimators=100, random_state=42)
    }

    results = {}

    for name, model in models.items():
        # 创建管道（包含标准化）
        pipeline = Pipeline([
            ('scaler', StandardScaler()),
            ('model', model)
        ])

        # 训练模型
        pipeline.fit(X_train, y_train)

        # 预测
        y_pred_train = pipeline.predict(X_train)
        y_pred_test = pipeline.predict(X_test)

        # 评估
        train_r2 = r2_score(y_train, y_pred_train)
        test_r2 = r2_score(y_test, y_pred_test)
        test_mse = mean_squared_error(y_test, y_pred_test)
        test_mae = mean_absolute_error(y_test, y_pred_test)

        results[name] = {
            'model': pipeline,
            'train_r2': train_r2,
            'test_r2': test_r2,
            'test_mse': test_mse,
            'test_mae': test_mae,
            'y_test': y_test,
            'y_pred_test': y_pred_test
        }

        print(f"\n{name} 结果:")
        print(f"训练集 R²: {train_r2:.4f}")
        print(f"测试集 R²: {test_r2:.4f}")
        print(f"测试集 MSE: {test_mse:.4f}")
        print(f"测试集 MAE: {test_mae:.4f}")

    return results, X_test, y_test


def 可视化结果(results):
    """
    可视化拟合结果
    """
    n_models = len(results)
    fig, axes = plt.subplots(2, 2, figsize=(15, 12))
    axes = axes.ravel()

    for idx, (name, result) in enumerate(results.items()):
        if idx < 4:  # 只显示前4个模型
            ax = axes[idx]
            ax.scatter(result['y_test'], result['y_pred_test'], alpha=0.6)
            ax.plot([result['y_test'].min(), result['y_test'].max()],
                    [result['y_test'].min(), result['y_test'].max()], 'r--', lw=2)
            ax.set_xlabel('实际值')
            ax.set_ylabel('预测值')
            # 修改标题，避免使用上标符号
            ax.set_title(f'{name}\nR2 = {result["test_r2"]:.4f}')
            ax.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.show()


def 特征重要性分析(results, feature_columns):
    """
    分析特征重要性（针对随机森林模型）
    """
    rf_result = results.get('Random Forest')
    if rf_result:
        model = rf_result['model']
        # 获取随机森林模型（在管道中的第二个步骤）
        rf_model = model.named_steps['model']

        importance = rf_model.feature_importances_

        # 创建特征重要性图
        plt.figure(figsize=(10, 6))
        indices = np.argsort(importance)[::-1]

        plt.bar(range(len(importance)), importance[indices])
        plt.xticks(range(len(importance)), [feature_columns[i] for i in indices], rotation=45)
        plt.title('特征重要性 (随机森林)')
        plt.tight_layout()
        plt.show()

        print("\n特征重要性排序:")
        for i in indices:
            print(f"{feature_columns[i]}: {importance[i]:.4f}")


def 保存模型(model, feature_columns, model_name):
    """
    保存训练好的模型到pkl文件
    """
    model_info = {
        'model': model,
        'feature_columns': feature_columns,
        'model_name': model_name
    }

    # 保存模型
    filename = f'背压预测模型_{model_name}.pkl'
    with open(filename, 'wb') as f:
        pickle.dump(model_info, f)

    print(f"\n模型已保存为: {filename}")
    print(f"特征列: {feature_columns}")


def 加载模型(filename):
    """
    从pkl文件加载模型
    """
    with open(filename, 'rb') as f:
        model_info = pickle.load(f)

    model = model_info['model']
    feature_columns = model_info['feature_columns']
    model_name = model_info['model_name']

    print(f"模型已从 {filename} 加载")
    print(f"模型名称: {model_name}")
    print(f"特征列: {feature_columns}")

    return model, feature_columns, model_name


def 预测新数据(model, new_data, feature_columns):
    """
    使用训练好的模型预测新数据
    """
    if isinstance(new_data, list):
        new_data = dict(zip(feature_columns, new_data))
    # 确保new_data是正确的格式
    if isinstance(new_data, dict):
        # 将字典转换为DataFrame
        import pandas as pd
        new_data = pd.DataFrame([new_data])

    # 确保列的顺序正确
    new_data = new_data[feature_columns]

    # 预测
    prediction = model.predict(new_data)
    return prediction


def 曲线拟合主函数():
    """
    主函数：执行完整的曲线拟合流程
    """
    # 1. 加载数据
    print("正在加载数据...")
    df = 加载工况整理后的数据()
    print(f"数据形状: {df.shape}")
    print(f"列名: {df.columns}")

    # 2. 数据预处理
    print("\n正在进行数据预处理...")
    X, y, feature_columns = 数据预处理(df)
    print(f"特征数量: {X.shape[1]}, 样本数量: {X.shape[0]}")

    # 3. 模型拟合
    print("\n正在训练多个模型...")
    results, X_test, y_test = 多种模型拟合(X, y, feature_columns)

    # 4. 选择最佳模型
    best_model_name = max(results.keys(), key=lambda k: results[k]['test_r2'])
    print(f"\n最佳模型: {best_model_name} (R² = {results[best_model_name]['test_r2']:.4f})")
    # 保存最佳模型
    best_model = results[best_model_name]['model']
    feature_columns_used = ['汽轮发电机功率', '汽机抽汽供热量', '环境温度', '大气湿度',
                            '循泵运行方式数字编码', '风机运行台数', '冷却塔进水蝶阀开启数量']

    保存模型(best_model, feature_columns_used, best_model_name)

    # 5. 可视化结果
    print("\n正在生成可视化结果...")
    可视化结果(results)

    # 6. 特征重要性分析
    特征重要性分析(results, feature_columns)

    return results, best_model_name


def 预测数据():
    model, feature_columns, model_name = 加载模型('背压预测模型_Random Forest.pkl')
    result = 预测新数据(model, [81.5, -17, 28, 12, 3, 4, 5], feature_columns)
    print(f"预测结果: {result}")


def 检查重复参数不同风机数量():
    """
    检查df1中是否存在前4个参数相近，循泵运行方式数字编码和冷却塔进水蝶阀开启数量相同，但风机运行台数不同的工况
    前4个参数：汽轮发电机功率、汽机抽汽供热量、环境温度、大气湿度（允许相近）
    相同参数：循泵运行方式数字编码、冷却塔进水蝶阀开启数量（要求完全相同）
    不同参数：风机运行台数（要求不同）
    """
    df1 = pl.read_csv("抽凝.csv")

    # 定义参数列
    continuous_params = [  # 连续参数，允许相近
        '汽轮发电机功率',
        '汽机抽汽供热量',
        '环境温度',
        '大气湿度'
    ]

    same_params = [  # 相同参数
        '循泵运行方式数字编码',
        '冷却塔进水蝶阀开启数量'
    ]

    different_param = '风机运行台数'  # 要求不同的参数

    # 检查列是否存在
    all_params = continuous_params + same_params + [different_param]
    missing_features = [col for col in all_params if col not in df1.columns]
    if missing_features:
        print(f"以下列在纯凝.csv中不存在: {missing_features}")
        return None

    # 为连续参数创建分组标签（将相近值归为一组）
    df_with_groups = df1.with_columns([
        # 汽轮发电机功率四舍五入到最接近的0.5
        (pl.col('汽轮发电机功率') * 2).round().alias('汽轮发电机功率_分组'),
        # 汽机抽汽供热量四舍五入到整数
        pl.col('汽机抽汽供热量').round().alias('汽机抽汽供热量_分组'),
        # 环境温度四舍五入到整数
        pl.col('环境温度').round().alias('环境温度_分组'),
        # 大气湿度四舍五入到整数
        pl.col('大气湿度').round().alias('大气湿度_分组')
    ])

    # 分组列：前4个参数的分组 + 2个相同参数
    group_columns = [
                        '汽轮发电机功率_分组',
                        '汽机抽汽供热量_分组',
                        '环境温度_分组',
                        '大气湿度_分组'
                    ] + same_params

    # 按分组列分组，查看每组中'风机运行台数'的值
    grouped = df_with_groups.group_by(group_columns).agg([
        pl.col(different_param).unique().sort().alias('不同风机台数'),
        pl.col(different_param).count().alias('记录数'),
        # 保留原始参数的范围以供参考
        pl.col('汽轮发电机功率').unique().sort().alias('汽轮发电机功率范围'),
        pl.col('汽机抽汽供热量').unique().sort().alias('汽机抽汽供热量范围'),
        pl.col('环境温度').unique().sort().alias('环境温度范围'),
        pl.col('大气湿度').unique().sort().alias('大气湿度范围'),
        # 保留示例值
        pl.first('循泵运行方式数字编码').alias('循泵运行方式数字编码_示例'),
        pl.first('冷却塔进水蝶阀开启数量').alias('冷却塔进水蝶阀开启数量_示例')
    ])

    # 找出有多个不同风机台数的组合
    different_fan_counts = grouped.filter(
        pl.col('不同风机台数').list.len() > 1
    )

    if len(different_fan_counts) > 0:
        print(f"发现 {len(different_fan_counts)} 种参数组合具有不同的风机运行台数:")
        for row in different_fan_counts.iter_rows(named=True):
            print(f"相近参数:")
            print(f"  汽轮发电机功率范围: {row['汽轮发电机功率范围']}")
            print(f"  汽机抽汽供热量范围: {row['汽机抽汽供热量范围']}")
            print(f"  环境温度范围: {row['环境温度范围']}")
            print(f"  大气湿度范围: {row['大气湿度范围']}")
            print(f"相同参数:")
            print(f"  循泵运行方式数字编码: {row['循泵运行方式数字编码_示例']}")
            print(f"  冷却塔进水蝶阀开启数量: {row['冷却塔进水蝶阀开启数量_示例']}")
            print(f"  不同风机台数: {row['不同风机台数']}")
            print(f"  记录数: {row['记录数']}")

            # 查找原始数据行
            conditions = [
                (pl.col('汽轮发电机功率_分组') == row['汽轮发电机功率_分组']),
                (pl.col('汽机抽汽供热量_分组') == row['汽机抽汽供热量_分组']),
                (pl.col('环境温度_分组') == row['环境温度_分组']),
                (pl.col('大气湿度_分组') == row['大气湿度_分组']),
                (pl.col('循泵运行方式数字编码') == row['循泵运行方式数字编码_示例']),
                (pl.col('冷却塔进水蝶阀开启数量') == row['冷却塔进水蝶阀开启数量_示例'])
            ]

            # 为每个条件添加
            filter_expr = conditions[0]
            for condition in conditions[1:]:
                filter_expr = filter_expr & condition

            matching_rows = df_with_groups.filter(filter_expr)

            print(f"  原始数据行:")
            for i, record in enumerate(matching_rows.iter_rows(named=True)):
                print(f"    记录{i + 1}:")
                # 输出所有列的值
                for col_name in df1.columns:
                    print(f"      {col_name}: {record[col_name]}")
                print()
    else:
        print("未发现前4个参数相近、循泵运行方式数字编码和冷却塔进水蝶阀开启数量相同但风机运行台数不同的工况")

    return different_fan_counts


def 构造新的数据集():
    df = pl.read_csv("抽凝.csv")
    for 汽轮发电机功率 in range(60, 290, 10):
        for 汽机抽汽供热量 in range(-30, 1910, 20):
            for 环境温度 in range(-10, 43, 2):
                for 大气湿度 in range(5, 100, 10):
                    for 循泵运行方式数字编码 in [1, 2, 3, 4, 5, 6]:
                        for 风机运行台数 in range(0, 9):
                            for 冷却塔进水蝶阀开启数量 in range(0, 9):
                                new_sample = {
                                    '汽轮发电机功率': 汽轮发电机功率,
                                    '汽机抽汽供热量': 汽机抽汽供热量,
                                    '环境温度': 环境温度,
                                    '大气湿度': 大气湿度,
                                    '循泵运行方式数字编码': 循泵运行方式数字编码,
                                    '风机运行台数': 风机运行台数,
                                    '冷却塔进水蝶阀开启数量': 冷却塔进水蝶阀开启数量
                                }
                                #


if __name__ == "__main__":
    # results, best_model_name = 曲线拟合主函数()
    # best_model = results[best_model_name]['model']
    # 预测数据()
    # 检查重复参数不同风机数量()
    构造新的数据集()
# 示例：预测新数据
# new_sample = {
#     '汽轮发电机功率': 100,
#     '汽机抽汽供热量': 50,
#     '环境温度': 20,
#     '大气湿度': 60,
#     '循泵运行方式数字编码': 1,
#     '风机运行台数': 2,
#     '冷却塔进水蝶阀开启数量': 3
# }
# prediction = 预测新数据(best_model, new_sample, ['汽轮发电机功率', '汽机抽汽供热量', '环境温度', '大气湿度', '循泵运行方式数字编码', '风机运行台数', '冷却塔进水蝶阀开启数量'])
# print(f"新数据预测结果: {prediction}")
