import polars as pl
from pathlib import Path
from yangke.common.config import logger
from yangke.base import plot_distribution


def 分段保存Excel文件(data, 文件名前缀, 分段大小=20000, folder=Path('D:/')):
    """
    将DataFrame分段保存到多个Excel文件中

    Parameters:
    data: polars DataFrame - 要保存的数据
    文件名前缀: str - 保存文件的前缀名称
    分段大小: int - 每个文件的最大行数，默认20000
    """
    total_rows = data.height
    num_chunks = (total_rows + 分段大小 - 1) // 分段大小

    logger.debug(f"数据总行数: {total_rows}, 将分割为 {num_chunks} 个文件")

    # 分段保存
    for i in range(num_chunks):
        start_row = i * 分段大小
        end_row = min((i + 1) * 分段大小, total_rows)
        chunk_df = data.slice(start_row, end_row - start_row)

        folder.mkdir(parents=True, exist_ok=True)  # 确保目录存在
        filename = folder / f"{文件名前缀}_第{i + 1}部分.xlsx"
        chunk_df.write_excel(filename)
        logger.debug(f"已保存 {filename.name} (行数: {chunk_df.height})")


def 过滤数据(data, filter_list=None, how='any'):
    """
    删除数据集中符合filter_list中任意条件的行

    Parameters:
    data: polars DataFrame
    filter_list: 过滤条件列表，如 ['汽轮发电机功率<60', '风速>20']
    how: 过滤方式，any表示任意条件满足就删除，all表示所有条件满足才删除

    Returns:
    过滤后的DataFrame
    """
    if filter_list is None:
        return data

    # 输出过滤前行数
    original_rows = data.height
    logger.debug(f"过滤前行数: {original_rows}")

    # 创建过滤条件
    filter_conditions = []

    for condition in filter_list:
        # 解析条件，例如 '汽轮发电机功率<60'
        if '<' in condition and '<=' not in condition:
            col_name, threshold = condition.split('<')
            threshold = float(threshold)
            filter_conditions.append(pl.col(col_name.strip()) < threshold)
        elif '>' in condition and '>=' not in condition:
            col_name, threshold = condition.split('>')
            threshold = float(threshold)
            filter_conditions.append(pl.col(col_name.strip()) > threshold)
        elif '<=' in condition:
            col_name, threshold = condition.split('<=')
            threshold = float(threshold)
            filter_conditions.append(pl.col(col_name.strip()) <= threshold)
        elif '>=' in condition:
            col_name, threshold = condition.split('>=')
            threshold = float(threshold)
            filter_conditions.append(pl.col(col_name.strip()) >= threshold)
        elif '==' in condition:
            col_name, threshold = condition.split('==')
            threshold = float(threshold)
            filter_conditions.append(pl.col(col_name.strip()) == threshold)
        elif '!=' in condition:
            col_name, threshold = condition.split('!=')
            threshold = float(threshold)
            filter_conditions.append(pl.col(col_name.strip()) != threshold)

    # 使用指定的逻辑组合条件
    if filter_conditions:
        if how == 'any':
            # any表示任意条件满足就删除
            combined_condition = filter_conditions[0]
            for cond in filter_conditions[1:]:
                combined_condition = combined_condition | cond

            # ~表示取反，即保留不满足任一条件的行
            data = data.filter(~combined_condition)
        elif how == 'all':
            # all表示所有条件都满足才删除
            combined_condition = filter_conditions[0]
            for cond in filter_conditions[1:]:
                combined_condition = combined_condition & cond

            # ~表示取反，即保留不满足所有条件的行（也就是至少有一个条件不满足的行）
            data = data.filter(~combined_condition)
        else:
            raise ValueError("how参数必须是'any'或'all'")

    # 输出过滤后行数
    filtered_rows = data.height
    logger.debug(
        f"过滤后行数{filter_list} : {filtered_rows}，删除行数: {original_rows - filtered_rows}，删除比例: {(original_rows - filtered_rows) / original_rows:.2%}")
    logger.debug(f"")

    return data


def 计算大泵运行台数(df_inner):
    """
    根据1号大循泵电流和2号大循泵电流计算大泵运行台数
    电流>200表示泵在运行，电流<5表示泵没有运行
    """
    # 计算大泵运行台数（电流大于200的泵数量）
    pump_running_conditions = [
        pl.col("1号大循泵电流（额定298A）") > 5,
        pl.col("2号大循泵电流（额定298A）") > 5
    ]

    df_inner = df_inner.with_columns(
        pl.sum_horizontal(pump_running_conditions).alias("大泵运行台数")
    )
    return df_inner


def 计算小泵运行台数(df_inner):
    """
    根据3号、4号小循泵电流计算小泵运行台数
    高速泵：3号小循泵电流（高速）和4号小循泵电流（高速）
    低速泵：3号小循泵电流（低速）和4号小循泵电流（低速）
    电流<5表示停机，否则运行
    """
    # 计算高速小泵运行台数
    high_speed_pump_conditions = [
        pl.col("3号小循泵电流（高速）") >= 5,
        pl.col("4号小循泵电流（高速）") >= 5
    ]

    df_inner = df_inner.with_columns(
        pl.sum_horizontal(high_speed_pump_conditions).alias("高速小泵运行台数")
    )

    # 计算低速小泵运行台数
    low_speed_pump_conditions = [
        pl.col("3号小循泵电流（低速）") >= 5,
        pl.col("4号小循泵电流（低速）") >= 5
    ]

    df_inner = df_inner.with_columns(
        pl.sum_horizontal(low_speed_pump_conditions).alias("低速小泵运行台数")
    )

    df_inner = df_inner.with_columns(
        (pl.col("高速小泵运行台数") + pl.col("低速小泵运行台数")).alias("小泵运行台数")
    )
    #
    # # 删除原始电流列
    # df_inner = df_inner.drop([
    #     "3号小循泵电流（高速）", "4号小循泵电流（高速）",
    #     "3号小循泵电流（低速）", "4号小循泵电流（低速）", "小泵运行台数"
    # ])

    return df_inner


def 计算风机运行台数(df_inner):
    # 计算风机运行台数
    """
    根据 机力塔风机1电流、机力塔风机2电流、机力塔风机3电流、机力塔风机4电流、机力塔风机5电流、机力塔风机6电流、机力塔风机7电流、机力塔风机8电流
    计算风机运行台数，风机电流大于20认为风机运行，风机电流小于5认为风机，删除电流在5和20之间的电流
    """
    fan_current_columns = [
        "机力塔风机1电流", "机力塔风机2电流", "机力塔风机3电流", "机力塔风机4电流",
        "机力塔风机5电流", "机力塔风机6电流", "机力塔风机7电流", "机力塔风机8电流"
    ]
    #
    # # 过滤掉风机电流在5-20之间的数据
    # fan_filter_conditions = []
    # for col in fan_current_columns:
    #     fan_filter_conditions.append((pl.col(col) >= 2) & (pl.col(col) <= 20))
    # if fan_filter_conditions:
    #     # 使用or操作符组合所有条件，过滤掉满足任一条件的行
    #     combined_fan_condition = fan_filter_conditions[0]
    #     for cond in fan_filter_conditions[1:]:
    #         combined_fan_condition = combined_fan_condition | cond
    #
    #     # ~表示取反，即保留不满足任一条件的行
    #     df_inner = df_inner.filter(~combined_fan_condition)
    #     print(f"已过滤掉风机电流在5-20之间的数据，过滤后行数: {df_inner.height}")
    # 计算风机运行台数（电流大于20的风机数量）
    fan_running_conditions = []
    for col in fan_current_columns:
        fan_running_conditions.append(pl.col(col) > 2)

    if fan_running_conditions:
        # 计算运行的风机台数
        df_inner = df_inner.with_columns(
            pl.sum_horizontal(fan_running_conditions).alias("风机运行台数")
        )

    # df_inner = df_inner.drop(fan_current_columns)
    return df_inner


def 标记运行台数变化后15分钟内数据(df_inner):
    """
    标记设备运行台数变化后的15分钟内数据
    新增一列 '删除标记'，包含删除原因信息
    """
    # 按时间排序
    df_inner = df_inner.sort("DateTime")

    # 获取所有运行台数列
    running_columns = ["风机运行台数", "大泵运行台数", "高速小泵运行台数", "低速小泵运行台数", "小泵运行台数"]

    # 初始化删除标记列为"不删除"
    df_inner = df_inner.with_columns(
        pl.lit("不删除").alias("删除标记")
    )

    # 为每个参数变化创建标记
    for col in running_columns:
        if col in df_inner.columns:
            # 计算该列的变化点（当前值不等于前一个值）
            change_mask = pl.col(col) != pl.col(col).shift(1)

            # 对于每个变化点，标记后续15行（15分钟）需要删除
            for i in range(0, 15):  # 1到15分钟
                row_mask = change_mask.shift(i).fill_null(False)  # 修改：使用正向偏移
                delete_reason = f"需删除，{col}变化"

                # 更新删除标记列
                df_inner = df_inner.with_columns(
                    pl.when((pl.col("删除标记") == "不删除") & row_mask)
                    .then(pl.lit(delete_reason))
                    .otherwise(pl.col("删除标记"))
                    .alias("删除标记")
                )

    # 统计删除行数
    rows_to_delete = df_inner.filter(pl.col("删除标记") != "不删除").height
    print(f"标记运行台数变化后15分钟内需要删除的数据: {rows_to_delete} 行")

    # 显示各种删除原因的统计
    reason_stats = df_inner.group_by("删除标记").agg(pl.count()).filter(pl.col("删除标记") != "不删除")
    print("删除原因统计:")
    for row in reason_stats.iter_rows():
        print(f"  {row[0]}: {row[1]} 行")

    return df_inner


def 标记数值快速变化列(df_inner, col_name="汽轮发电机功率", time_interval=10, val_change=5, next_minutes=5):
    """
    标记指定列在指定时间间隔内数值变化超过阈值的行，并将之后指定时间的数据也标记为删除

    Parameters:
    df_inner: DataFrame
    col_name: 要检查的列名，默认为"汽轮发电机功率"
    time_interval: 时间间隔（分钟），默认为10分钟
    val_change: 数值变化阈值，默认为5
    next_minutes: 变化后需要标记的分钟数，默认为5分钟
    """
    # 按时间排序
    df_inner = df_inner.sort("DateTime")

    # 初始化删除标记列（如果不存在）
    if "删除标记" not in df_inner.columns:
        df_inner = df_inner.with_columns(
            pl.lit("不删除").alias("删除标记")
        )

    # 检查指定列是否存在
    if col_name in df_inner.columns:
        # 计算指定时间间隔前的数值（time_interval行前，因为数据间隔为1分钟）
        val_interval_ago = pl.col(col_name).shift(time_interval)

        # 计算数值变化绝对值
        value_change = (pl.col(col_name) - val_interval_ago).abs()

        # 标记数值变化超过阈值的点
        rapid_change_mask = value_change > val_change

        # 对于每个快速变化点，标记该点及之后next_minutes分钟的数据
        for i in range(0, next_minutes + 1):  # 0到next_minutes分钟
            row_mask = rapid_change_mask.shift(i).fill_null(False)
            delete_reason = f"需删除，{col_name}快速变化"

            # 更新删除标记列
            df_inner = df_inner.with_columns(
                pl.when((pl.col("删除标记") == "不删除") & row_mask)
                .then(pl.lit(delete_reason))
                .otherwise(pl.col("删除标记"))
                .alias("删除标记")
            )
    else:
        logger.debug(f"警告：列 '{col_name}' 不存在于数据中")

    # 统计删除行数
    rows_to_delete = df_inner.filter(pl.col("删除标记") != "不删除").height
    logger.debug(f"标记{col_name}快速变化需要删除的数据: {rows_to_delete} 行")

    return df_inner


def filter_data(df_inner, folder):
    """
    根据特定规则过滤数据，这里主要是指定规则

    """
    # 输出内存占用
    memory_usage = df_inner.estimated_size()
    logger.debug(f"数据内存占用: {memory_usage / (1024 * 1024):.2f} MB")

    df_inner = 计算风机运行台数(df_inner)
    df_inner = 计算大泵运行台数(df_inner)
    df_inner = 计算小泵运行台数(df_inner)
    df_inner = 标记运行台数变化后15分钟内数据(df_inner)
    df_inner = 标记数值快速变化列(df_inner, col_name="汽轮发电机功率", time_interval=10, val_change=5, next_minutes=5)

    # df_inner = 标记数值快速变化列(df_inner, col_name="瞬时热网供水热量", time_interval=10, val_change=5, next_minutes=5)
    df_inner = df_inner.with_columns(
        (pl.col("大泵运行台数").cast(pl.Utf8) +
         pl.col("高速小泵运行台数").cast(pl.Utf8) +
         pl.col("低速小泵运行台数").cast(pl.Utf8))
        .alias("循泵运行方式")
    )

    # 定义循泵运行方式映射字典
    循泵运行方式映射 = {
        "000": 0,
        "001": 1,  # 1台低速小泵
        "010": 2,  # 1台高速小泵
        "100": 3,  # 1台大泵
        "020": 4,  # 2台高速小泵
        "110": 5,  # 一大一小高速
        "200": 6,  # 2台大泵
        "120": 7,  # 一大两小高速（不常用） # 经测试后删除7和8这两种方式，可以有效提高神经网络预测准确度
        "210": 8,  # 两大一小高速（不常用）
    }

    # 使用polars的map_elements方法进行转换
    df_inner = df_inner.with_columns(
        pl.col("循泵运行方式")
        .map_elements(
            lambda x: 循泵运行方式映射.get(x, -1),
            return_dtype=pl.Int32
        )
        .alias("循泵运行方式数字编码")
    )
    # plot_distribution(df_inner, "循泵运行方式数字编码", "循泵运行方式分布")

    # 分段保存合并后的数据
    _folder = folder / '合并原始数据据表'
    _folder.mkdir(parents=True, exist_ok=True)  # 确保目录存在
    # 分段保存Excel文件(df_inner, "合并后的历史数据", 2000, _folder)

    # 删除标记为需要删除的行
    original_rows = df_inner.height
    df_inner = df_inner.filter(pl.col("删除标记") == "不删除")
    deleted_rows = original_rows - df_inner.height
    print(f"删除标记为需要删除的行: {deleted_rows} 行")

    df_inner = 过滤数据(df_inner, ['循泵运行方式数字编码<=0'])
    df_inner = 过滤数据(df_inner, ['循泵运行方式数字编码>=7'])

    df_inner = df_inner.with_columns(
        ((pl.col("大气压力1") + pl.col("大气压力2")) / 20000)
        .alias("平均大气压力")  # 转换单位为MPa
    )
    # plot_distribution(df_inner, "平均大气压力", "平均大气压力分布") # 绘制数据分布图

    df_inner = df_inner.with_columns(
        ((pl.col("#1燃机环境温度") + pl.col("#2燃机环境温度")) / 2)
        .alias("环境温度")
    )
    df_inner = df_inner.with_columns(
        ((pl.col("#1燃机进气湿度") + pl.col("#2燃机进气湿度")) / 2)
        .alias("大气湿度") # 湿度范围0~100
    )

    # plot_distribution(df_inner, "汽轮发电机功率", "汽轮发电机功率分布")
    df_inner = 过滤数据(df_inner, ['汽轮发电机功率<60'])
    # 删除总功率<145MW 的数据
    df_inner = df_inner.with_columns(
        (pl.col("燃机功率1") + pl.col("燃机功率2") + pl.col("汽轮发电机功率"))
        .alias("机组总功率")
    )
    df_inner = 过滤数据(df_inner, ['机组总功率<145'])

    # plot_distribution(df_inner, "风速", "风速分布")
    # df_inner = 过滤数据(df_inner, ['平均大气压力>0.105', '平均大气压力<0.095'])

    # plot_distribution(df_inner, "凝结水补水流量", "补水流量分布")
    df_inner = 过滤数据(df_inner, ['凝结水补水流量>50'])

    # df_inner = 过滤数据(df_inner, ['低压缸入口调门开度<5'])  # 低压缸不进汽，背压模式，不优化，直接指定开一台低速泵即可
    df_inner = 过滤数据(df_inner, ['背压模式>0.5'])  # 背压模式=1表示是背压模式运行
    df_inner = df_inner.with_columns(
        (pl.col("1号冷却塔进水蝶阀") + pl.col("2号冷却塔进水蝶阀") + pl.col("3号冷却塔进水蝶阀") + pl.col(
            "4号冷却塔进水蝶阀")
         + pl.col("5号冷却塔进水蝶阀")
         + pl.col("6号冷却塔进水蝶阀")
         + pl.col("7号冷却塔进水蝶阀")
         + pl.col("8号冷却塔进水蝶阀")
         )
        .alias("冷却塔进水蝶阀开启数量")
    )

    df_inner = 过滤数据(df_inner, ['AGC指令<0.5']) # AGC指令=0表示AGC退出了，此时可能

    df_inner = df_inner.with_columns(
        (pl.col("1号炉烟气供热器供热量") + pl.col("2号炉烟气供热器供热量"))
        .alias("烟气供热器供热量")
    )
    # plot_distribution(df_inner, "烟气供热器供热量", "烟气供热器供热量分布")

    df_inner = df_inner.with_columns(
        (pl.col("瞬时热网供水热量") - pl.col("瞬时热网回水热量"))
        .alias("总供热量")
    )
    df_inner = df_inner.with_columns(
        (pl.col("总供热量") - pl.col("烟气供热器供热量"))  # 烟气供热为凝结水再循环部分和烟气直接传递给热网循环水的热量，没有进入汽轮机系统，需要扣除
        .alias("汽机抽汽供热量")
    )
    # plot_distribution(df_inner, "汽机抽汽供热量", "汽机抽汽供热量分布")

    df_inner = df_inner.with_columns(
        (pl.col("1号炉高旁减温水流量") + pl.col("2号炉高旁减温水流量"))
        .alias("高旁减温水流量")
    )
    # plot_distribution(df_inner, "高旁减温水流量", "高旁减温水流量分布")

    df_inner = df_inner.with_columns(
        (pl.col("汽机循环水供水管道流量1") + pl.col("汽机循环水供水管道流量2"))
        .alias("循环水流量")
    )

    # plot_distribution(df_inner, "循环水流量", "循环水流量分布")

    df_inner = df_inner.with_columns(
        ((pl.col("凝汽器真空1") + pl.col("凝汽器真空2")) / 2)
        .alias("背压")
    )

    df_inner = df_inner.with_columns(
        (pl.col("凝汽器真空1") - pl.col("凝汽器真空选择值")).abs()
        .alias("背压差1")
    )
    df_inner = df_inner.with_columns(
        (pl.col("凝汽器真空2") - pl.col("凝汽器真空选择值")).abs()
        .alias("背压差2")
    )
    df_inner = 过滤数据(df_inner, ['背压差1>0.3', '背压差2>0.3'])

    # df_inner = 过滤数据(df_inner, ['循泵运行方式数字编码==4', '循环水流量<31000'], how='all')
    # df_inner = 过滤数据(df_inner, ['循泵运行方式数字编码==5', '循环水流量<39000'], how='all')

    # 删除不再使用的数据列
    df_inner = df_inner.drop(
        ["大气压力1", "大气压力2",
         "1号炉烟气供热器供热量", "2号炉烟气供热器供热量", "烟气供热器供热量",
         "热网供水压力",
         "热网回水压力", "热网供水温度", "热网回水温度",
         '1号炉高旁减温水流量', '2号炉高旁减温水流量',
         "凝汽器真空1", "凝汽器真空2", "凝汽器真空选择值", "背压差1", "背压差2"
         ])

    # 输出过滤后内存占用
    memory_usage_filtered = df_inner.estimated_size()
    logger.debug(f"过滤后数据内存占用: {memory_usage_filtered / (1024 * 1024):.2f} MB")
    return df_inner
