import torch
import glob
import os
import pandas as pd
import sys
import datetime as dt
import numpy as np

from tqdm import tqdm
from torch import nn
from torch.nn import functional as F
from torch.utils.data import Dataset, DataLoader, Sampler
from datetime import datetime, timedelta

from schedule import *



class custom_dataset(Dataset):

    def __init__(self, folder_paths=None,
                 years=None, past_len=6, future_len=6, interval=3, dtype_=torch.float32):
        """
        Args:
            folder_paths: 数据文件夹路径
            years: 年份列表，格式为 [year1, year2, ...]，如 [2020, 2021, 2022]
            past_len: 过去时间步长
            future_len: 未来时间步长
            interval: 时间间隔
            dtype_: 数据类型
        """
        self.dtype_ = dtype_

        # 合并所有年份的时间窗口
        all_window_list = {
            "past_times": [],
            "future_times": []
        }

        for year in years:
            # 为每个年份生成整年的时间窗口
            start_time = datetime(year, 1, 1, 1, 30)
            end_time = datetime(year, 12, 31, 22, 30)

            window_list = generate_hour_windows(
                start_date=start_time,
                end_date=end_time,
                past_len=past_len,
                future_len=future_len,
                interval=interval
            )
            # 合并时间窗口
            all_window_list["past_times"].extend(window_list["past_times"])
            all_window_list["future_times"].extend(window_list["future_times"])

        filter_window_list = []
        past_times = []
        future_times = []

        for x, y in zip(all_window_list["past_times"], all_window_list["future_times"]):

            valid = True
            x_files = []
            y_files = []

            # Check if all past time steps exist
            for x_ in x:
                year, month, day, hour, minute = convert_datetime_to_str(x_)
                file_path = self.return_exists_file_path(folder_paths, f"{year}{month}{day}{hour}{minute}.pth")
                # file_path = f"{folder_paths}/{year}{month}{day}{hour}{minute}.pth"
                if not os.path.exists(file_path):
                    valid = False
                    break
                x_files.append(file_path)

            # Check if all future time steps exist
            if valid:
                for y_ in y:
                    year, month, day, hour, minute = convert_datetime_to_str(y_)
                    file_path = self.return_exists_file_path(folder_paths, f"{year}{month}{day}{hour}{minute}.pth")
                    # file_path = f"{folder_paths}/{year}{month}{day}{hour}{minute}.pth"
                    if not os.path.exists(file_path):
                        valid = False
                        break
                    y_files.append(file_path)

            # If all time steps exist, add the window and file paths to the filter list
            if valid:
                filter_window_list.append((x, y, x_files, y_files))
                past_times.append(x)
                future_times.append(y)

        # self.filter_window_list = filter_window_list[:100]
        self.filter_window_list = filter_window_list
        self.past_times = past_times
        self.future_times = future_times

        print(f"总窗口数量为{len(self.filter_window_list)}")

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

    def __getitem__(self, idx):

        _, _, x_files, y_files = self.filter_window_list[idx]
        x_info = torch.load(x_files[0])
        y_info = torch.load(y_files[0])

        time = torch.tensor([x_info["hour"] / 24, x_info["dayofweek"] / 7, x_info["dayofyear"] / 365.25])
        x = x_info["data"]
        y = y_info["data"]

        # Exclude channels DUSMASS25 (index 12) and SSSMASS25 (index 13)
        x = torch.cat([x[:12], x[14:]], dim=0)
        y = torch.cat([y[:12], y[14:]], dim=0)

        return x.to(self.dtype_), y.to(self.dtype_), time.to(self.dtype_)

    def return_exists_file_path(self, folder_list, file_name):

        found = False
        for folder in folder_list:
            filepath = os.path.join(folder, file_name)
            if os.path.exists(filepath):
                found = True
                return filepath
        if found == False:
            return "xxx"