# data_preprocessor.py
import os
import torch
import joblib
import numpy as np
import pandas as pd
from tqdm import tqdm    # 进度条显示
from sklearn.preprocessing import MinMaxScaler    # 数据归一化工具
from torch.utils.data import DataLoader, TensorDataset    # PyTorch数据加载工具
from concurrent.futures import ThreadPoolExecutor    # 多线程读取文件

class DataPreprocessor:
    """数据预处理类，负责数据加载、划分、转换为模型可输入的格式"""
    
    @staticmethod
    def load_and_process_data(args, data):
        
        """
        加载并处理数据，划分训练/验证/测试集，创建数据加载器
        参数：
            args: 配置参数（包含数据集划分日期、序列长度等）
            data: 预处理后的完整数据（含日期列）
        返回：
            train_loader: 训练集数据加载器
            val_loader: 验证集数据加载器
            test_loader: 测试集数据加载器
            data: 原始数据（用于后续处理）
        """
        
        # 处理日期列
        data['date'] = pd.to_datetime(data['date'])
        time_interval = pd.Timedelta(minutes=(4 * args.resolution / 60))
        window_time_span = time_interval * (args.seq_len + 1)

        # 划分训练/验证/测试集（调整起始日期以适应滑动窗口）
        val_start_date = pd.to_datetime(args.val_start_date)
        test_start_date = pd.to_datetime(args.test_start_date)
        
        # 调整验证集/测试集起始时间（向前推一个窗口，确保有足够历史数据构建输入序列）
        adjusted_val_start = val_start_date - window_time_span
        adjusted_test_start = test_start_date - window_time_span
        
        # 构建数据集掩码（按日期筛选）
        train_mask = (data['date'] >= pd.to_datetime(args.train_start_date)) & \
                     (data['date'] <= pd.to_datetime(args.train_end_date))

        val_mask = (data['date'] >= adjusted_val_start) & \
                   (data['date'] <= pd.to_datetime(args.val_end_date))

        test_mask = (data['date'] >= adjusted_test_start) & \
                    (data['date'] <= pd.to_datetime(args.test_end_date))

        # 筛选数据并重置索引
        train_data = data[train_mask].reset_index(drop=True)
        val_data = data[val_mask].reset_index(drop=True)
        test_data = data[test_mask].reset_index(drop=True)
        
        # 移除日期列用于建模
        train_data = train_data.drop(columns=['date'])
        val_data = val_data.drop(columns=['date'])
        test_data = test_data.drop(columns=['date'])
    
        # 创建监督学习数据集（输入序列+目标序列）
        train_supervised = DataPreprocessor.create_supervised_dataset(
            args,
            train_data,
            1
        )
        
        val_supervised = DataPreprocessor.create_supervised_dataset(
            args,
            val_data,
            1
        )
        
        test_supervised = DataPreprocessor.create_supervised_dataset(
            args,
            test_data,
            args.step_size
        )
        
        # 转换为DataLoader
        train_loader = DataPreprocessor.load_data(
            args, 
            train_supervised,
            shuffle=True
        )
        
        val_loader = DataPreprocessor.load_data(
            args, 
            val_supervised,
            shuffle=False
        )
        
        test_loader = DataPreprocessor.load_data(
            args, 
            test_supervised,
            shuffle=False
        )
        
        return train_loader, val_loader, test_loader, data  # 返回原始数据用于后续处理
    
    @staticmethod
    def read_and_combine_csv_files(args):
        """
        多线程读取并合并多个CSV文件，进行下采样、日期处理和归一化
        参数：
            args: 配置参数（包含数据路径、文件范围等）
        返回：
            chunk: 预处理后的合并数据（含日期和归一化特征）
        """
        current_dir = os.path.dirname(__file__)
        parent_dir = os.path.dirname(current_dir)
        args.data_dir = os.path.join(parent_dir, args.data_dir)
        
        def read_file(file_count):
            """读取单个CSV文件的函数（供多线程调用）"""
            file_name = args.file_pattern.format(file_count)
            file_path = os.path.join(args.data_dir, file_name)
            return pd.read_csv(file_path)
        
        # 生成待读取的文件索引列表
        file_indices = list(range(args.start_files, args.end_files + 1))
        
        # 多线程读取文件（加速大文件读取）
        max_workers = os.cpu_count()
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            results = list(tqdm(executor.map(read_file, file_indices),
                                total=len(file_indices),
                                desc="正在读取文件"))
        
        all_data = pd.concat(results, ignore_index=True)
        
        # 根据feature_num筛选特征列
        if args.feature_num == 16:
            # 定义需要保留的16个特征（包含index列用于后续日期处理）
            '''
            specified_features = [
                "C.M.RO1_FT_JS@out",
                "C.M.RO2_FT_JS@out",
                "C.M.RO3_FT_JS@out",
                "C.M.RO4_FT_JS@out",
                "C.M.RO_TT_ZJS@out",
                "C.M.RO_Cond_ZCS@out",
                "C.M.RO1_DB@DPT_1",
                "C.M.RO1_DB@DPT_2",
                "C.M.RO2_DB@DPT_1",
                "C.M.RO2_DB@DPT_2",
                "C.M.RO3_DB@DPT_1",
                "C.M.RO3_DB@DPT_2",
                "C.M.RO4_DB@DPT_1",
                "C.M.RO4_DB@DPT_2"
            ]
            '''
            specified_features = [
                "C.M.UF1_FT_JS@out",        # UF1进水流量
                "C.M.UF2_FT_JS@out",        # UF2进水流量
                "C.M.UF3_FT_JS@out",        # UF3进水流量
                "C.M.UF4_FT_JS@out",        # UF4进水流量
                "C.M.RO_TT_ZJS@out",        # 反渗透总进水温度
                "C.M.UF_ORP_ZCS@out",       # 超滤总产水ORP
                "C.M.UF1_DB@press_PV",      # UF1跨膜压差
                "C.M.UF2_DB@press_PV",      # UF2跨膜压差
                "C.M.UF3_DB@press_PV",      # UF3跨膜压差
                "C.M.UF4_DB@press_PV",      # UF4跨膜压差
                "UF1Per",                   # UF1膜渗透率
                "UF2Per",                   # UF2膜渗透率
                "UF3Per",                   # UF3膜渗透率
                "UF4Per",                   # UF4膜渗透率
                ]

            # 必须保留'index'列用于后续日期处理，添加到特征列表
            columns_to_keep = ['index'] + specified_features
            # 筛选并按指定顺序保留列
            all_data = all_data[columns_to_keep]
        # 当feature_num=79时，保持原有所有特征
        
        # 按分辨率下采样
        chunk = all_data.iloc[::args.resolution, :].reset_index(drop=True)
        
        # 处理日期和时间特征
        chunk = DataPreprocessor.process_date(chunk, args)
        # 归一化
        chunk = DataPreprocessor.scaler_data(chunk, args)
        
        return chunk
    
    @staticmethod
    def process_date(data, args):
        """
        处理日期列，根据分辨率生成周期性时间特征
        参数：
            data: 含'index'列（原始日期）的DataFrame
            resolution: 数据分辨率，用于决定生成的时间特征
        返回：
            data: 处理后的DataFrame（含日期列和时间特征）
        """
        data = data.rename(columns={'index': 'date'})
        data['date'] = pd.to_datetime(data['date'])
    
        # 生成周期性时间特征
        time_features = []
        
        if args.resolution == 60:
            # 分辨率为60时，生成分钟级和日级特征
            data['minute_of_day'] = data['date'].dt.hour * 60 + data['date'].dt.minute
            data['minute_sin'] = np.sin(2 * np.pi * data['minute_of_day'] / 1440)
            data['minute_cos'] = np.cos(2 * np.pi * data['minute_of_day'] / 1440)
            time_features.extend(['minute_sin', 'minute_cos'])
            data.drop(columns=['minute_of_day'], inplace=True)
        
        # 两种分辨率下都保留日级特征
        data['day_of_year'] = data['date'].dt.dayofyear
        data['day_year_sin'] = np.sin(2 * np.pi * data['day_of_year'] / 366)
        data['day_year_cos'] = np.cos(2 * np.pi * data['day_of_year'] / 366)
        time_features.extend(['day_year_sin', 'day_year_cos'])
        data.drop(columns=['day_of_year'], inplace=True)
    
        # 调整列顺序（日期+时间特征+其他特征）
        other_columns = [col for col in data.columns if col not in ['date'] and col not in time_features]
        data = data[['date'] + time_features + other_columns]
    
        return data
    
    @staticmethod
    def scaler_data(data, args):
        """
        对数据进行归一化（0-1缩放），并保存归一化器（供预测时反归一化）
        参数：
            data: 含'date'列和特征列的DataFrame
            args: 配置参数（包含scaler_path）
        返回：
            scaled_data: 归一化后的DataFrame（含日期列）
        """
        date_col = data[['date']]
        data_to_scale = data.drop(columns=['date'])

        scaler = MinMaxScaler(feature_range=(0, 1))
        scaled_data = scaler.fit_transform(data_to_scale)
        joblib.dump(scaler, args.scaler_path)  # 保存归一化器

        # 转换为DataFrame并拼接日期列
        scaled_data = pd.DataFrame(scaled_data, columns=data_to_scale.columns)
        scaled_data = pd.concat([date_col.reset_index(drop=True), scaled_data], axis=1)
        
        return scaled_data
    
    @staticmethod
    def create_supervised_dataset(args, data, step_size):
        """
        创建监督学习数据集（输入序列+目标序列）
        输入序列：历史seq_len个时间步的所有特征
        目标序列：未来output_size个时间步的标签特征（最后labels_num列）
        参数：
            args: 配置参数（含seq_len、output_size等）
            data: 输入数据（不含日期列的特征数据）
            step_size: 采样步长（每隔step_size取一个样本）
        返回：
            dataset: 监督学习数据集（DataFrame）
        """
        data = pd.DataFrame(data)
        cols = []
        col_names = []
        
        feature_columns = data.columns.tolist()

        # 输入序列（t-0到t-(seq_len-1)）
        for col in feature_columns:
            for i in range(args.seq_len - 1, -1, -1):
                cols.append(data[[col]].shift(i))
                col_names.append(f"{col}(t-{i})")
        
        # 目标序列（仅取最后labels_num列作为预测目标）
        target_columns = feature_columns[-args.labels_num:]
        for i in range(1, args.output_size + 1):
            for col in target_columns:
                cols.append(data[[col]].shift(-i))
                col_names.append(f"{col}(t+{i})")

        # 合并并清洗数据
        dataset = pd.concat(cols, axis=1)
        dataset.columns = col_names
        dataset = dataset.iloc[::step_size, :]  # 按步长采样
        dataset.dropna(inplace=True)  # 移除含缺失值的行
        
        return dataset

    @staticmethod
    def load_data(args, dataset, shuffle):
        """
        将监督学习数据集转换为PyTorch张量，并创建DataLoader
        参数：
            args: 配置参数（含特征数、批大小等）
            dataset: 监督学习数据集（DataFrame）
            shuffle: 是否打乱数据（训练集True，验证/测试集False）
        返回：
            data_loader: PyTorch DataLoader
        """
        input_length = args.seq_len
        n_features = args.feature_num
        labels_num = args.labels_num
    
        n_features_total = n_features * input_length  # 输入特征总维度
        n_labels_total = args.output_size * labels_num  # 目标总维度

        # 分割输入和目标
        X = dataset.values[:, :n_features_total]
        y = dataset.values[:, n_features_total:n_features_total + n_labels_total]
    
        # 重塑输入为[样本数, 序列长度, 特征数]
        X = X.reshape(X.shape[0], input_length, n_features)
        X = torch.tensor(X, dtype=torch.float32).to(args.device)
        y = torch.tensor(y, dtype=torch.float32).to(args.device)

        # 创建数据集和数据加载器
        dataset_tensor = TensorDataset(X, y)
        generator = torch.Generator()
        generator.manual_seed(args.random_seed)  # 固定随机种子确保可复现
        
        data_loader = DataLoader(
            dataset_tensor, 
            batch_size=args.batch_size, 
            shuffle=shuffle,
            generator=generator
        )
    
        return data_loader