# -*- coding: utf-8 -*-
# @Author: Jin HuaiMing
# @Date: 2025-08-11 15:19
from torch.utils.data import Dataset,DataLoader
import pandas as pd
import torch
from datetime import datetime, timedelta
import os
import numpy as np
from sklearn.preprocessing import StandardScaler
import joblib
import glob

class PipelineSolarDataset(Dataset):
    def __init__(self, 
                input_folder,
                features,
                target,
                start_time,
                end_time,
                pre_seq,
                proj = 'OL_JIANGMEN',
                model_name = 'Dlinear',
                release = 'v1',
                saved_model_path = 'None',
                mode = 'train',
                capacity = 4050
    ):
        
        self.features = features
        self.target = target
        self.pre_seq = pre_seq
        self.proj = proj
        self.model_name = model_name
        self.release = release
        self.saved_model_path = saved_model_path
        self.mode = mode
        self.capacity = capacity

        if self.release == 'None':
            self.out_path = os.path.join(self.saved_model_path,self.proj,self.model_name,str(self.pre_seq))
        else:
            self.out_path = os.path.join(self.saved_model_path,self.proj,self.model_name,str(self.pre_seq),self.release)
        
        os.makedirs(self.out_path, exist_ok=True)
        files = []
        start_time = datetime.strptime(str(start_time), "%Y%m%d")
        end_time = datetime.strptime(str(end_time), "%Y%m%d")

        """按天遍历日期范围"""
        current = start_time
        while current <= end_time:
            file_00 = f'{input_folder}/{self.proj}/{current.strftime("%Y%m%d")}00.csv'
            file_12 = f'{input_folder}/{self.proj}/{current.strftime("%Y%m%d")}12.csv'

            if os.path.exists(file_00):
                files.append(file_00)
            if os.path.exists(file_12):
                files.append(file_12)
            current += timedelta(days=1)
        
        
        self.data_list = []
        # 将files每个文件对应的features列拼接到data_Scaler组成一个大的数组 (Sample, feature)
        scaler_arrays = []
        for f in files:
            df = pd.read_csv(f)
            # 判断截取pre_seq行，长度等于pre_seq，否则跳过
            if len(df.iloc[:self.pre_seq]) < self.pre_seq:
                continue

            # 保存原有data_list格式，便于原有get_item用
            self.data_list.append(
                {
                    "x": df[self.features].to_numpy(),
                    "y": df[self.target].to_numpy(),
                    'capacity': 4050
                }
            )

            # 训练模型制作归一化器
            if self.mode  == 'train':
                # 只收集features, 拼接
                scaler_arrays.append(df[self.features].to_numpy()[:self.pre_seq])

                # 训练模型制作归一化器
                if self.mode  == 'train':
                    # 只收集features, 拼接
                    scaler_arrays.append(df[self.features].to_numpy()[:self.pre_seq])

        if self.mode == 'train':
            scaler_arrays = np.concatenate(scaler_arrays, axis=0)
            # 保存归一化器到本地（例如.pkl文件，放到当前目录下）
            self.scaler = StandardScaler()
            self.scaler.fit(scaler_arrays)
            scaler_path = os.path.join(self.out_path, "scaler.pkl")
            joblib.dump(self.scaler, scaler_path)
        else:
            scaler_path = os.path.join(self.out_path, "scaler.pkl")
            if not os.path.exists(scaler_path):
                raise FileNotFoundError(f"Scaler file not found: {self.proj}-{self.model_name}-{self.pre_seq}-{self.release}")
            self.scaler = joblib.load(scaler_path)
 
    def __getitem__(self, idx):
        d = self.data_list[idx]
        # 应用self.scaler对x,y做归一化
        x_scaled = self.scaler.transform(d["x"][0:self.pre_seq, :])
        x = torch.tensor(x_scaled, dtype=torch.float32).view(self.pre_seq, d["x"].shape[1])
        y = torch.tensor(d["y"][0:self.pre_seq] / d['capacity'], dtype=torch.float32)

        # x = [batch, input_len, n_features]
        # y = [batch, output_len]
        return (x, y)

    def __len__(self):
        return len(self.data_list)

if __name__ == "__main__":
    
    data = PipelineSolarDataset(
        "/home/jinhuaiming/OuLiPowerV1.0/data/short_term",
        ['ssrd','tcc','lcc','mcc','hcc','2t','2d','10u','10v'],
        '实际发电功率（kW）',
        "20250614",
        "20260331",
        288,
        proj = 'JMJD01',
        model_name = 'GRU',
        release = 'v24',
        saved_model_path = '/home/jinhuaiming/OuLiPowerV1.0/saved_model',
        mode = 'train',
        capacity = 4050
    )
    
    dataload = DataLoader(dataset=data, batch_size=1)
    print(len(data))
    for item in dataload:
        x1, x2 = item
        print(x2)