import polars as pl


def 统计参数分布(df: pl.DataFrame, col1='#3发电机有功功率1', col2='环境温度1', col1_list=None,
                 col2_list=None, transpose=False, output_file="result.xlsx"):
    """
    df中存在power_col和temp_col两列，分别表示功率和环境温度，根据power_list和temp_list对数据进行分组统计。
    统计每个功率区间和每个环境温度区间内的样本数量。包括区间两侧的样本区间，例如[0, 10]表示统计3个区间的样本数量，分别是≤0, (0, 10], >10。
    输出结果是个表格，行标题表示功率区间，列标题表示环境温度区间，每个单元格表示该功率区间和环境温度区间内的样本数量。如：

	运行时间/min <170	170~238	238~272	272~306	>306
    <0	        34503	2467	1329	1279	1352
	0~10	    93909	5169	5990	6091	8069
	10~20	    129031	4219	6302	4644	10923
	>20	        18155	755	    2328	1379	6170

    :param df:
    :param col1:
    :param col2:
    :param col1_list:
    :param col2_list:
    :param transpose: 是否转置结果表格，默认False
    :param output_file:
    :return:
    """
    if col1_list is None:
        col1_list = [170, 238, 272, 306]
    if col2_list is None:
        col2_list = [0, 10, 20]

    def create_bins(breaks_list):
        """根据输入的列表创建完整的分箱边界，避免使用inf"""
        if not breaks_list:
            return []
        # 使用极大极小值代替无穷大
        min_val = float('-1e308')  # 接近负无穷的值
        max_val = float('1e308')  # 接近正无穷的值
        bins = [min_val]
        bins.extend(sorted(breaks_list))
        bins.append(max_val)
        return bins

    def create_labels(breaks_list, label_type):
        """根据输入的列表创建区间标签"""
        if not breaks_list:
            return []

        breaks_sorted = sorted(breaks_list)
        labels = []

        # 第一个区间: < first_break
        labels.append(f"<{breaks_sorted[0]}")

        # 中间区间: [start, end)
        for i in range(len(breaks_sorted) - 1):
            labels.append(f"{breaks_sorted[i]}~{breaks_sorted[i + 1]}")

        # 最后一个区间: >= last_break
        labels.append(f">={breaks_sorted[-1]}")  # 改为>=避免包含边界点的问题

        return labels

    def manual_binning(df, power_col, temp_col, power_bins, temp_bins, power_labels, temp_labels):
        """手动实现分箱逻辑 - 完全使用polars"""

        # 定义分配函数
        def assign_to_bins(series, bins, labels):
            """为polars Series分配区间标签"""
            if len(bins) == 0 or len(labels) == 0:
                return pl.lit(None).cast(pl.Utf8)

            # 从第一个条件开始构建表达式
            result_expr = pl.lit(None).cast(pl.Utf8)

            # 从最后一个区间向前构建条件链，这样更高效
            for i in range(len(bins) - 1):
                lower_bound = bins[i]
                upper_bound = bins[i + 1]

                # 构建条件：lower_bound <= value < upper_bound
                condition = (series >= lower_bound) & (series < upper_bound)

                # 使用when-then-otherwise构建条件链
                result_expr = pl.when(condition).then(labels[i]).otherwise(result_expr)

            return result_expr

        # 为功率和温度分配区间
        df_with_bins = df.with_columns([
            assign_to_bins(pl.col(power_col), power_bins, power_labels).alias("功率区间"),
            assign_to_bins(pl.col(temp_col), temp_bins, temp_labels).alias("温度区间")
        ])

        # 过滤掉None值的行
        df_with_bins = df_with_bins.filter(
            (pl.col("功率区间").is_not_null()) &
            (pl.col("温度区间").is_not_null())
        )

        # 统计交叉表
        result = df_with_bins.group_by(["功率区间", "温度区间"]).len().pivot(
            values="len",
            index="功率区间",
            on="温度区间"  # 使用新的参数名 'on' 替代 'columns'
        )

        return result

    # 定义一个排序函数，用于解析区间字符串并按数值排序
    def parse_interval(interval_str):
        if isinstance(interval_str, str) and interval_str.startswith('<'):
            # 形如 "<170" 的区间，取数字部分
            number_part = interval_str[1:].lstrip('=').strip()  # 去除可能的等号和空格
            if number_part:  # 确保数字部分不为空
                return (0, float(number_part))
            else:
                return (0, float('-inf'))  # 空数字部分视为负无穷
        elif isinstance(interval_str, str) and '~' in interval_str:
            # 形如 "170~238" 的区间，取起始值
            start_val = interval_str.split('~')[0].strip()
            if start_val:  # 确保数字部分不为空
                return (1, float(start_val))
            else:
                return (1, 0.0)  # 默认值
        elif isinstance(interval_str, str) and interval_str.startswith('>='):
            # 形如 ">=306" 的区间，取数字部分
            number_part = interval_str[2:].strip()
            if number_part:  # 确保数字部分不为空
                return (2, float(number_part))
            else:
                return (2, float('inf'))  # 空数字部分视为正无穷
        elif isinstance(interval_str, str) and interval_str.startswith('>'):
            # 形如 ">306" 的区间，取数字部分
            number_part = interval_str[1:].strip()
            if number_part:  # 确保数字部分不为空
                return (2, float(number_part))
            else:
                return (2, float('inf'))  # 空数字部分视为正无穷
        elif isinstance(interval_str, str) and interval_str.startswith('<='):
            # 形如 "<=306" 的区间，取数字部分
            number_part = interval_str[2:].strip()
            if number_part:  # 确保数字部分不为空
                return (0, float(number_part))
            else:
                return (0, float('-inf'))  # 空数字部分视为负无穷
        elif isinstance(interval_str, str) and interval_str.startswith('=='):
            # 形如 "==306" 的区间，取数字部分
            number_part = interval_str[2:].strip()
            if number_part:  # 确保数字部分不为空
                return (1, float(number_part))
            else:
                return (1, 0.0)  # 默认值
        else:
            # 其他情况返回原值
            return (3, interval_str if isinstance(interval_str, (int, float)) else 0.0)

    # 创建功率区间标签
    power_bins = create_bins(col1_list)
    power_labels = create_labels(col1_list, "功率")

    # 创建环境温度区间标签
    temp_bins = create_bins(col2_list)
    temp_labels = create_labels(col2_list, "温度")

    # 尝试使用Polars的cut函数进行分箱
    try:
        # 为功率创建区间
        df_result = df.with_columns([
            pl.col(col1).cut(breaks=power_bins[1:-1], labels=power_labels).alias("功率区间").cast(pl.Utf8),
            # 移除首尾的极值
            pl.col(col2).cut(breaks=temp_bins[1:-1], labels=temp_labels).alias("温度区间").cast(pl.Utf8)  # 移除首尾的极值
        ])

        # 统计每个组合的计数
        result = df_result.group_by(["功率区间", "温度区间"]).len().pivot(
            values="len",
            index="功率区间",
            on="温度区间"  # 使用新的参数名 'on' 替代 'columns'
        )

    except Exception as e:
        print(f"Error during binning: {e}")
        # 如果cut函数不可用，则手动实现分箱逻辑
        result = manual_binning(df, col1, col2, power_bins, temp_bins, power_labels, temp_labels)

    # 如果需要转置表格
    if transpose:
        # 获取当前的列名（温度区间）
        current_cols = [col for col in result.columns if col != '功率区间']

        # 转置表格：将原来的行变为列，列变为行
        # 首先将表格重塑为长格式，使用 unpivot 替代 melt
        melted = result.unpivot(
            on=current_cols,  # 使用新的参数名 'on' 替代 'value_vars'
            index=['功率区间'],  # 使用新的参数名 'index' 替代 'id_vars'
            variable_name='温度区间',
            value_name='样本数量'
        )

        # 然后以温度区间为行，功率区间为列进行透视
        transposed_result = melted.pivot(
            values='样本数量',
            index='温度区间',
            on='功率区间'  # 使用新的参数名 'on' 替代 'columns'
        )

        # 对转置后的表格进行排序
        # 获取当前的列名（功率区间）
        power_cols = [col for col in transposed_result.columns if col != '温度区间']

        # 对温度区间进行排序
        unique_temp_intervals = transposed_result.select('温度区间').unique()
        temp_values = []
        for row in unique_temp_intervals.iter_rows():
            temp_values.append(row[0])
        temp_order = sorted(temp_values, key=parse_interval)

        ordered_temp_expr = pl.when(pl.col('温度区间') == temp_order[0]).then(0)
        for i, temp_interval in enumerate(temp_order[1:], 1):
            ordered_temp_expr = ordered_temp_expr.when(pl.col('温度区间') == temp_interval).then(i)
        ordered_temp_expr = ordered_temp_expr.otherwise(len(temp_order))

        # 对功率区间进行排序
        power_order = sorted(power_cols, key=parse_interval)

        # 创建带有排序索引的结果
        result_with_order = transposed_result.with_columns(ordered_temp_expr.alias('_temp_order'))

        # 重新排序并删除辅助列
        result_sorted = result_with_order.sort('_temp_order').drop('_temp_order')

        # 重新排列列顺序（功率区间）
        final_result = result_sorted.select(['温度区间'] + power_order)
        final_result.write_excel(output_file)
        return final_result
    else:
        # 按照数值大小对功率区间和温度区间进行排序
        # 获取当前的列名（温度区间）
        temp_cols = [col for col in result.columns if col != '功率区间']

        # 对功率区间进行排序
        if '功率区间' in result.columns:
            unique_power_intervals = result.select('功率区间').unique()
            power_values = []
            for row in unique_power_intervals.iter_rows():
                power_values.append(row[0])

            power_order = sorted(power_values, key=parse_interval)
            ordered_power_expr = pl.when(pl.col('功率区间') == power_order[0]).then(0)
            for i, power_interval in enumerate(power_order[1:], 1):
                ordered_power_expr = ordered_power_expr.when(pl.col('功率区间') == power_interval).then(i)
            ordered_power_expr = ordered_power_expr.otherwise(len(power_order))

            # 对温度区间进行排序
            temp_order = sorted(temp_cols, key=parse_interval)

            # 创建带有排序索引的结果
            result_with_order = result.with_columns(ordered_power_expr.alias('_power_order'))

            # 重新排序并删除辅助列
            result_sorted = result_with_order.sort('_power_order').drop('_power_order')

            # 重新排列列顺序（温度区间）
            final_result = result_sorted.select(['功率区间'] + temp_order)
            final_result.write_excel(output_file)
            return final_result
        result.write_excel(output_file)
        return result


if __name__ == "__main__":
    df = pl.read_excel(r"/home/yke/Documents/2025/临河数据/9-12月份数据.xlsx")
    print("\n历史运行时间统计:")
    print(统计参数分布(df, transpose=True, output_file="表1.xlsx"))

    print("\n冷端系统不同运行方式时间占比（功率：0~170MW）:")
    print(统计参数分布(df.filter(pl.col('#3发电机有功功率1') <= 170), col1='环境温度1', col2='平均风机频率',
                       col1_list=[0, 10, 20], col2_list=[16, 20, 30, 40], output_file="表2.xlsx"))

    print("\n冷端系统不同运行方式时间占比（功率：170~238MW）:")
    print(统计参数分布(df.filter((pl.col('#3发电机有功功率1') > 170) & (pl.col('#3发电机有功功率1') <= 238)),
                       col1='环境温度1', col2='平均风机频率',
                       col1_list=[0, 10, 20], col2_list=[16, 20, 30, 40],  output_file="表3.xlsx"))

    print("\n冷端系统不同运行方式时间占比（功率：238~272MW）:")
    print(统计参数分布(df.filter((pl.col('#3发电机有功功率1') > 238) & (pl.col('#3发电机有功功率1') <= 272)),
                       col1='环境温度1', col2='平均风机频率',
                       col1_list=[0, 10, 20], col2_list=[16, 20, 30, 40], output_file="表4.xlsx"))

    print("\n冷端系统不同运行方式时间占比（功率：272~306MW）:")
    print(统计参数分布(df.filter((pl.col('#3发电机有功功率1') > 272) & (pl.col('#3发电机有功功率1') <= 306)),
                       col1='环境温度1', col2='平均风机频率',
                       col1_list=[0, 10, 20], col2_list=[16, 20, 30, 40], output_file="表5.xlsx"))

    print("\n冷端系统不同运行方式时间占比（功率：>306MW）:")
    print(统计参数分布(df.filter(pl.col('#3发电机有功功率1') > 306), col1='环境温度1', col2='平均风机频率',
                       col1_list=[0, 10, 20], col2_list=[16, 20, 30, 40], output_file="表6.xlsx"))
