from yangke.common.fileOperate import read_and_merge_data
import polars as pl
from yangke.base import interpolate_null_values
import numpy as np

# folder = r'C:\Users\54067\WPSDrive\217145378\WPS云盘\7性能试验\2025\国电投临河电厂\空冷数据\空冷优化'
# folder = r'C:\Users\YangKe\WPSDrive\217145378_3\WPS云盘\7性能试验\2025\国电投临河电厂\空冷数据\空冷优化'
# folder = r'C:\Users\wysun\WPSDrive\217145378_1\WPS云盘\7性能试验\2025\国电投临河电厂\空冷数据\空冷优化'
folder = r'/home/yke/Documents/2025/临河数据'
files = ['20250901_14.xlsx', '20250901_24.xlsx', '20250901_34.xlsx']
# files = ['20251001_13.xlsx', '20251001_23.xlsx', '20251001_33.xlsx']
# files = ['20251101_12.xlsx', '20251101_22.xlsx', '20251101_32.xlsx']
# files = ['20251201_11.xlsx', '20251201_21.xlsx', '20251201_31.xlsx']


def save_selected_period(df: pl.DataFrame, start_time, end_time):
    """
    按照时间列DateTime，将选定的时间段的数据保存到xlsx文件中
    """
    # 第一步，首先判断DateTime列和start_time和end_time是否是datetime类型，如果不是，则转换成datetime类型
    if df.schema["DateTime"] != pl.Datetime:
        df = df.with_columns(pl.col("DateTime").str.strptime(pl.Datetime, "%Y-%m-%d %H:%M:%S"))

    # 确保start_time和end_time是datetime类型
    if not isinstance(start_time, pl.Datetime):
        if isinstance(start_time, str):
            start_time = pl.Series([start_time]).str.strptime(pl.Datetime, "%Y-%m-%d %H:%M:%S")[0]
        else:
            start_time = pl.Series([start_time]).cast(pl.Datetime)[0]

    if not isinstance(end_time, pl.Datetime):
        if isinstance(end_time, str):
            end_time = pl.Series([end_time]).str.strptime(pl.Datetime, "%Y-%m-%d %H:%M:%S")[0]
        else:
            end_time = pl.Series([end_time]).cast(pl.Datetime)[0]

    # 第二步，从df中筛选出处于start_time和end_time之间的数据
    filtered_df = df.filter(
        (pl.col("DateTime") >= start_time) &
        (pl.col("DateTime") <= end_time)
    )

    # 第三步，保存到xlsx文件中
    output_file = f"selected_period_{start_time.strftime('%Y%m%d_%H%M%S')}_{end_time.strftime('%Y%m%d_%H%M%S')}.xlsx"
    filtered_df.write_excel(output_file)

    print(f"已保存选定时间段的数据到文件: {output_file}")
    print(f"共保存 {len(filtered_df)} 行数据")

    return filtered_df


def find_specific_condition(df: pl.DataFrame, conditions: dict):
    """
       从df中查找符合指定条件的数据的时间段，返回满足指定条件的时间段的开始时间和结束时间，时间列名为DateTime。
       查询条件为：
       {
           "#3发电机有功功率1": "between 339 and 341",
           "环境温度1": "vibration<1", # 表示时间段内，参数波动小于1
           "排汽母管压力1": "vibration<1",
           "3号机冷再至工业供汽管道流量计算后": "vibration<1",
           "机主蒸汽抽汽双减后流量计算后": "vibration<1"
       }
       """
    # 首先确保DateTime列是datetime类型
    if df.schema["DateTime"] != pl.Datetime:
        df = df.with_columns(pl.col("DateTime").str.strptime(pl.Datetime, "%Y-%m-%d %H:%M:%S"))

    # 用于标记满足所有条件的行
    condition_mask = pl.lit(True)

    for column, condition in conditions.items():
        if "vibration" in condition:
            # 处理振动条件，即参数波动小于指定值
            threshold = float(condition.split("<")[1])

            # 计算该列的滚动统计，找到满足波动条件的连续时间段
            # 首先我们需要计算滑动窗口内的波动（标准差或极差）
            series_data = df[column]

            # 创建一个函数来检查滑动窗口是否满足振动条件
            def check_vibration_condition(window_data, threshold):
                # 计算窗口内数据的波动（标准差）
                std_dev = np.std(window_data)
                return std_dev < threshold

            # 创建一个标记列，标记满足振动条件的行
            # 由于振动条件涉及连续数据段，我们需要查找满足条件的连续时间段
            data_values = df[column].to_numpy()
            vibration_mask = pl.Series([False] * len(df))

            # 逐个检查可能的连续时间段
            for i in range(len(data_values)):
                # 检查从i开始的子序列，直到不满足振动条件
                for j in range(i + 2, len(data_values) + 1):  # 至少需要2个点才能计算波动
                    segment = data_values[i:j]
                    if np.std(segment) >= threshold:
                        break
                    vibration_mask[j - 1] = True  # 标记该点满足振动条件

            # 创建一个列，表示满足振动条件的行
            vibration_col = pl.Series(f"{column}_vibration_ok", vibration_mask)
            condition_mask = condition_mask & pl.lit(vibration_col)
        else:
            # 处理范围条件，如 "between 339 and 341"
            if "between" in condition:
                parts = condition.replace("between", "").replace("and", " ").split()
                min_val = float(parts[0])
                max_val = float(parts[1])

                # 添加范围条件到条件掩码
                condition_mask = condition_mask & (pl.col(column).is_between(min_val, max_val))

    # 应用所有条件筛选数据
    filtered_df = df.filter(condition_mask)

    if len(filtered_df) == 0:
        print("没有找到满足条件的数据段")
        return None, None

    # 找到满足条件的时间段的开始和结束时间
    start_time = filtered_df["DateTime"].min()
    end_time = filtered_df["DateTime"].max()

    print(f"找到满足条件的时间段: {start_time} 到 {end_time}，共 {len(filtered_df)} 行数据")

    return start_time, end_time


def find_stable_period(df: pl.DataFrame, stability_thresholds: dict, min_length: int = 10):
    """
    查找整个时间段内极差满足阈值的所有连续稳定段。
    参数：
    - df: 输入的DataFrame
    - stability_thresholds: 每列的极差阈值
    - min_length: 最小时间段长度（行数）
    返回：
    - [(start_time, end_time), ...]
    """

    def check_stability(data_, columns_, i_, j_):
        for col in columns_:
            segment_ = data_[col][i_:j_ + 1]
            max_val_ = np.max(segment_)
            min_val_ = np.min(segment_)
            if max_val_ - min_val_ > stability_thresholds[col]:
                return False
        return True

    if df.schema["DateTime"] != pl.Datetime:
        df = df.with_columns(pl.col("DateTime").str.strptime(pl.Datetime, "%Y-%m-%d %H:%M:%S"))

    df = df.sort("DateTime")
    columns = list(stability_thresholds.keys())
    n = len(df)

    stable_periods = []

    # 转换为numpy数组以便快速计算
    data = {col: df[col].to_numpy() for col in columns}
    times = df["DateTime"].to_list()

    # 遍历所有可能的时间段
    for i in range(n - min_length):
        j = i + min_length - 1
        valid = check_stability(data, columns, i, j) and j < n

        if valid:
            while check_stability(data, columns, i, j + 1) and j + 1 < n:
                j = j + 1  # 如果满足长度的时间段极差小于给定值，在逐步增长该时间段，直到不满足极差要求

            start_time = times[i]
            end_time = times[j]
            stable_periods.append((start_time, end_time))

            # 可选：打印该段的极差
            print(f"稳定时间段: {start_time} 到 {end_time}")
            for col in columns:
                segment = data[col][i:j + 1]
                max_val = np.max(segment)
                min_val = np.min(segment)
                print(f"  {col}: 极差 = {max_val - min_val:.3f}")
            print()
            i = j + 1  # 下一个时间段从已获取时间段的后面开始
        if i % 100 == 0:
            print(f"已处理 {i} 行数据")

    print(f"找到 {len(stable_periods)} 个满足全局极差条件的稳定时间段")
    return stable_periods


def fix_column_title(df: pl.DataFrame):
    """
    将列标题转换为正确的格式，如果列标题包含\n换行符，就删除换行符和换行符之前的内容，只保留后面的内容
    """
    # 获取所有列名
    old_columns = df.columns

    # 创建新的列名映射
    new_columns = {}
    for col in old_columns:
        if '\n' in col:
            # 如果列名包含换行符，保留换行符之后的内容
            new_col = col.split('\n')[-1]
            new_columns[col] = new_col
        else:
            # 如果列名不包含换行符，保持不变
            new_columns[col] = col

    # 重命名列
    for old_col, new_col in new_columns.items():
        if old_col != new_col:
            df = df.rename({old_col: new_col})

    return df


if __name__ == '__main__':
    merged_data: pl.DataFrame = read_and_merge_data(folder, files)
    # 将merged_data的数据类型转换为浮点数，有些是字符串
    merged_data = fix_column_title(merged_data)
    interpolated_data = interpolate_null_values(merged_data)
    interpolated_data.write_excel('9月份数据.xlsx')
    # interpolated_data.write_excel('10月份数据.xlsx')
    # interpolated_data.write_excel('11月份数据.xlsx')
    # interpolated_data.write_excel('12月份数据.xlsx')

    # interpolated_data = pl.read_excel('11月份数据.xlsx')
    # interpolated_data = pl.read_excel('12月份数据.xlsx')
    # stability_thresholds = {
    #     "3号机冷再至工业供汽管道流量计算后": 1,
    #     "#3发电机有功功率1": 1,
    #     "#3发电机有功功率2": 1,
    #     "#3发电机有功功率3": 1,
    #     "排汽母管压力1": 0.5,
    #     "排汽母管压力2": 0.5,
    #     "排汽母管压力3": 0.5,
    #     "环境温度1": 0.5,
    #     "环境温度2": 0.5,
    #     "环境温度3": 0.5,
    #     "第一列＃1风机频率": 0.5,
    #     "第一列＃2风机频率": 0.5,
    #     "第一列＃3风机频率": 0.5,
    #     "第一列＃4风机频率": 0.5,
    #     "第一列＃5风机频率": 0.5,
    #     "第二列＃1风机频率": 0.5,
    #     "第二列＃2风机频率": 0.5,
    #     "第二列＃3风机频率": 0.5,
    #     "第二列＃4风机频率": 0.5,
    #     "第二列＃5风机频率": 0.5,
    #     "第三列＃1风机频率": 0.5,
    #     "第三列＃2风机频率": 0.5,
    #     "第三列＃3风机频率": 0.5,
    #     "第三列＃4风机频率": 0.5,
    #     "第三列＃5风机频率": 0.5,
    #     "第四列＃1风机频率": 0.5,
    #     "第四列＃2风机频率": 0.5,
    #     "第四列＃3风机频率": 0.5,
    #     "第四列＃4风机频率": 0.5,
    #     "第四列＃5风机频率": 0.5,
    #     "第五列＃1风机频率": 0.5,
    #     "第五列＃2风机频率": 0.5,
    #     "第五列＃3风机频率": 0.5,
    #     "第五列＃4风机频率": 0.5,
    #     "第五列＃5风机频率": 0.5,
    #     "第六列＃1风机频率": 0.5,
    #     "第六列＃2风机频率": 0.5,
    #     "第六列＃3风机频率": 0.5,
    #     "第六列＃4风机频率": 0.5,
    #     "第六列＃5风机频率": 0.5,
    #     "机主蒸汽抽汽双减后流量计算后": 1
    # }
    # stable_periods = find_stable_period(interpolated_data, stability_thresholds, min_length=30)
    # for i, (start_time, end_time) in enumerate(stable_periods):
    #     print(f"稳定时间段: {start_time} 到 {end_time}")
    #     # 保存时间段数据到xlsx文件
    #     period_data = interpolated_data.filter(
    #         (pl.col("DateTime") >= start_time) & (pl.col("DateTime") <= end_time)
    #     )
    #     filename = f"{start_time.strftime('%Y%m%d%H%M%S')}.xlsx"
    #     period_data.write_excel(filename)
    #     print(f"时间段数据已保存到: {filename}")
