import os
import numpy as np
import pandas as pd
import os
import torch
from torch.utils.data import Dataset, DataLoader
from sklearn.preprocessing import StandardScaler
from utils.timefeatures import time_features
import warnings

warnings.filterwarnings('ignore')

class Dataset_MMD(Dataset):
    def __init__(self, args, root_path, flag='train', size=None,
                 features='S', traindata_path='synthetic_train_dataset.pth', valdata_path='synthetic_val_dataset.pth',
                 predict_path='synthetic_val_dataset.pth', target='OT', scale=True, timeenc=0, freq='h', seasonal_patterns=None, train_only=False):
        # size [seq_len, label_len, pred_len]
        self.args = args
        # info
        if size == None:
            self.seq_len = 24 * 4 * 4
            self.label_len = 24 * 4
            self.pred_len = 24 * 4
        else:
            self.seq_len = size[0]
            self.label_len = size[1]
            self.pred_len = size[2]
        # init
        assert flag in ['train', 'test', 'val', 'pred']
        type_map = {'train': 0, 'val': 1, 'test': 2, 'pred': 3}
        self.set_type = type_map[flag]

        self.features = features
        self.target = target
        self.scale = scale
        self.timeenc = timeenc
        self.freq = freq

        self.root_path = root_path
        self.traindata_path = traindata_path
        self.valdata_path = valdata_path
        self.predict_path = predict_path
        self.train_only = train_only
        self.__read_data__()
        self.tot_len = len(self.data_x)

    def __read_data__(self):
        self.scaler = StandardScaler()
        
        if self.set_type == 0:
            train_data = torch.load(os.path.join(self.root_path,
                                          self.traindata_path))
            train_data_x = []
            train_data_y = []
            train_data_text = []
            for i in train_data.keys():
                train_data_x.append(train_data[i]['ts_x'].numpy())
                train_data_y.append(train_data[i]['ts_y'].numpy())
                train_data_text.append(train_data[i]['history_text'])
            self.data_x = train_data_x
            self.data_y = train_data_y
            self.text = train_data_text
            
        elif self.set_type == 1:
            test_data = torch.load(os.path.join(self.root_path,
                                          self.valdata_path))
            test_data_x = []
            test_data_y = []
            test_data_text = []
            for i in test_data.keys():
                test_data_x.append(test_data[i]['ts_x'].numpy())
                test_data_y.append(test_data[i]['ts_y'].numpy())
                test_data_text.append(test_data[i]['history_text'])
            self.data_x = test_data_x
            self.data_y = test_data_y
            self.text = test_data_text
            
        elif self.set_type == 2:
            val_data = torch.load(os.path.join(self.root_path,
                                          self.valdata_path))
            val_data_x = []
            val_data_y = []
            val_data_text = []
            for i in val_data.keys():
                val_data_x.append(val_data[i]['ts_x'].numpy())
                val_data_y.append(val_data[i]['ts_y'].numpy())
                val_data_text.append(val_data[i]['history_text'])
            self.data_x = val_data_x
            self.data_y = val_data_y
            self.text = val_data_text
        if self.set_type == 3:
            pred_data = torch.load(os.path.join(self.root_path,
                                          self.predict_path))
            pred_data_x = []
            pred_data_y = []
            pred_data_text = []
            for i in pred_data.keys():
                pred_data_x.append(pred_data[i]['ts_x'].numpy())
                pred_data_y.append(pred_data[i]['ts_y'].numpy())
                pred_data_text.append(pred_data[i]['history_text'])
            self.data_x = pred_data_x
            self.data_y = pred_data_y
            self.text = pred_data_text
    def get_text(self, indices):
        if isinstance(indices, torch.Tensor):
            indices = indices.numpy()
        
        text=np.array([self.text[indice] for indice in indices])
        # 使用高级索引提取数据
        return text
    def __getitem__(self, index):
        seq_x = self.data_x[index]
        seq_y = self.data_y[index]
        seq_x_mark = torch.zeros((seq_x.shape[0], 1))
        seq_y_mark = torch.zeros((seq_x.shape[0], 1))

        return seq_x, seq_y, seq_x_mark, seq_y_mark,index

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

    def inverse_transform(self, data):
        return self.scaler.inverse_transform(data)


class Dataset_synthetic(Dataset):
    def __init__(self, args, root_path, flag='train', size=None,
                 features='S', traindata_path='synthetic_train_dataset.pth', valdata_path='synthetic_val_dataset.pth',
                 predict_path='synthetic_val_dataset.pth',target='OT', scale=True, timeenc=0, freq='h', train_only=False):
        # size [seq_len, label_len, pred_len]
        # info
        self.seq_len = size[0]
        self.label_len = size[1]
        self.pred_len = size[2]
        # init
        assert flag in ['train', 'test', 'val', 'pred']
        type_map = {'train': 0, 'val': 1, 'test': 2, 'pred': 3}
        self.set_type = type_map[flag]

        self.features = features
        self.target = target
        self.scale = scale
        self.timeenc = timeenc
        self.freq = freq

        self.root_path = root_path
        self.traindata_path = traindata_path
        self.valdata_path = valdata_path
        self.predict_path = predict_path
        self.__read_data__()

    def __read_data__(self):
        self.scaler = StandardScaler()     

        if self.set_type == 0:
            train_data = torch.load(os.path.join(self.root_path,
                                          self.traindata_path))
            train_data_x = []
            train_data_y = []
            for i in train_data.keys():
                train_data_x.append(train_data[i]['ts_x'].numpy())
                train_data_y.append(train_data[i]['ts_y'].numpy())
            self.data_x = train_data_x
            self.data_y = train_data_y
            
        elif self.set_type == 1:
            val_data = torch.load(os.path.join(self.root_path,
                                            self.valdata_path))
            val_data_x = []
            val_data_y = []
            for i in val_data.keys():
                val_data_x.append(val_data[i]['ts_x'].numpy())
                val_data_y.append(val_data[i]['ts_y'].numpy())
            self.data_x = val_data_x
            self.data_y = val_data_y
            
        elif self.set_type == 2:
            test_data = torch.load(os.path.join(self.root_path,
                                          self.valdata_path))
            test_data_x = []
            test_data_y = []
            for i in test_data.keys():
                test_data_x.append(test_data[i]['ts_x'].numpy())
                test_data_y.append(test_data[i]['ts_y'].numpy())
            self.data_x = test_data_x
            self.data_y = test_data_y
            
        elif self.set_type == 3:
            pred_data = torch.load(os.path.join(self.root_path,
                                            self.predict_path))
            pred_data_x = []
            pred_data_y = []
            for i in pred_data.keys():
                pred_data_x.append(pred_data[i]['ts_x'].numpy())
                pred_data_y.append(pred_data[i]['ts_y'].numpy())
            self.data_x = pred_data_x
            self.data_y = pred_data_y

    def __getitem__(self, index):

        seq_x = self.data_x[index]
        seq_y = self.data_y[index]
        seq_x_mark = torch.zeros((seq_x.shape[0], 1))
        seq_y_mark = torch.zeros((seq_x.shape[0], 1))

        return seq_x, seq_y, seq_x_mark, seq_y_mark

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

    def inverse_transform(self, data):
        return self.scaler.inverse_transform(data)

class Dataset_public(Dataset):
    def __init__(self, args, root_path, flag='train', size=None,
                 features='S', traindata_path='synthetic_train_dataset.pth', valdata_path='synthetic_val_dataset.pth',
                 predict_path='synthetic_val_dataset.pth',target='OT', scale=True, timeenc=0, freq='h', train_only=False):
        # size [seq_len, label_len, pred_len]
        # info
        self.seq_len = size[0]
        self.label_len = size[1]
        self.pred_len = size[2]
        # init
        assert flag in ['train', 'test', 'val', 'pred']
        type_map = {'train': 0, 'val': 1, 'test': 2, 'pred': 3}
        self.set_type = type_map[flag]

        self.features = features
        self.target = target
        self.scale = scale
        self.timeenc = timeenc
        self.freq = freq

        self.root_path = root_path
        self.traindata_path = traindata_path
        self.valdata_path = valdata_path
        self.predict_path = predict_path
        self.__read_data__()

    def __read_data__(self):
        self.scaler = StandardScaler()     

        if self.set_type == 0:
            train_data = torch.load(os.path.join(self.root_path,
                                          self.traindata_path))
            train_data_x = []
            train_data_y = []
            for i in train_data.keys():
                train_data_x.append(train_data[i]['ts_x'].numpy())
                train_data_y.append(train_data[i]['ts_y'].numpy())
            self.data_x = train_data_x
            self.data_y = train_data_y
            
        elif self.set_type == 1:
            val_data = torch.load(os.path.join(self.root_path,
                                            self.valdata_path))
            val_data_x = []
            val_data_y = []
            for i in val_data.keys():
                val_data_x.append(val_data[i]['ts_x'].numpy())
                val_data_y.append(val_data[i]['ts_y'].numpy())
            self.data_x = val_data_x
            self.data_y = val_data_y
            
        elif self.set_type == 2:
            test_data = torch.load(os.path.join(self.root_path,
                                          self.valdata_path))
            test_data_x = []
            test_data_y = []
            for i in test_data.keys():
                test_data_x.append(test_data[i]['ts_x'].numpy())
                test_data_y.append(test_data[i]['ts_y'].numpy())
            self.data_x = test_data_x
            self.data_y = test_data_y
            
        elif self.set_type == 3:
            pred_data = torch.load(os.path.join(self.root_path,
                                            self.predict_path))
            pred_data_x = []
            pred_data_y = []
            for i in pred_data.keys():
                pred_data_x.append(pred_data[i]['ts_x'].numpy())
                pred_data_y.append(pred_data[i]['ts_y'].numpy())
            self.data_x = pred_data_x
            self.data_y = pred_data_y

    def __getitem__(self, index):

        seq_x = self.data_x[index]
        seq_y = self.data_y[index]
        seq_x_mark = torch.zeros((seq_x.shape[0], 1))
        seq_y_mark = torch.zeros((seq_x.shape[0], 1))

        return seq_x, seq_y, seq_x_mark, seq_y_mark

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

    def inverse_transform(self, data):
        return self.scaler.inverse_transform(data)


class Dataset_Custom(Dataset):
    def __init__(self, root_path, flag='train', size=None,
                 features='S', data_path='ETTh1.csv',
                 target='OT', scale=True, timeenc=0, freq='h', train_only=False):
        # size [seq_len, label_len, pred_len]
        # info
        if size == None:
            self.seq_len = 24 * 4 * 4
            self.label_len = 24 * 4
            self.pred_len = 24 * 4
        else:
            self.seq_len = size[0]
            self.label_len = size[1]
            self.pred_len = size[2]
        # init
        assert flag in ['train', 'test', 'val']
        type_map = {'train': 0, 'val': 1, 'test': 2}
        self.set_type = type_map[flag]

        self.features = features
        self.target = target
        self.scale = scale
        self.timeenc = timeenc
        self.freq = freq
        self.train_only = train_only

        self.root_path = root_path
        self.data_path = data_path
        self.__read_data__()

    def __read_data__(self):
        self.scaler = StandardScaler()
        df_raw = pd.read_csv(os.path.join(self.root_path,
                                          self.data_path))
        '''
        df_raw.columns: ['date', ...(other features), target feature]
        '''
        cols = list(df_raw.columns)
        if self.features == 'S':
            cols.remove(self.target)
        cols.remove('date')
        # print(cols)
        num_train = int(len(df_raw) * (0.7 if not self.train_only else 1))
        num_test = int(len(df_raw) * 0.2)
        num_vali = len(df_raw) - num_train - num_test
        border1s = [0, num_train - self.seq_len, len(df_raw) - num_test - self.seq_len]
        border2s = [num_train, num_train + num_vali, len(df_raw)]
        border1 = border1s[self.set_type]
        border2 = border2s[self.set_type]

        if self.features == 'M' or self.features == 'MS':
            df_raw = df_raw[['date'] + cols]
            cols_data = df_raw.columns[1:]
            df_data = df_raw[cols_data]
        elif self.features == 'S':
            df_raw = df_raw[['date'] + cols + [self.target]]
            df_data = df_raw[[self.target]]

        if self.scale:
            train_data = df_data[border1s[0]:border2s[0]]
            self.scaler.fit(train_data.values)
            # print(self.scaler.mean_)
            # exit()
            data = self.scaler.transform(df_data.values)
        else:
            data = df_data.values

        df_stamp = df_raw[['date']][border1:border2]
        df_stamp['date'] = pd.to_datetime(df_stamp.date)
        if self.timeenc == 0:
            df_stamp['month'] = df_stamp.date.apply(lambda row: row.month, 1)
            df_stamp['day'] = df_stamp.date.apply(lambda row: row.day, 1)
            df_stamp['weekday'] = df_stamp.date.apply(lambda row: row.weekday(), 1)
            df_stamp['hour'] = df_stamp.date.apply(lambda row: row.hour, 1)
            data_stamp = df_stamp.drop(['date'], 1).values
        elif self.timeenc == 1:
            data_stamp = time_features(pd.to_datetime(df_stamp['date'].values), freq=self.freq)
            data_stamp = data_stamp.transpose(1, 0)

        self.data_x = data[border1:border2]
        self.data_y = data[border1:border2]
        self.data_stamp = data_stamp

    def __getitem__(self, index):
        s_begin = index
        s_end = s_begin + self.seq_len
        r_begin = s_end - self.label_len
        r_end = r_begin + self.label_len + self.pred_len

        seq_x = self.data_x[s_begin:s_end]
        seq_y = self.data_y[r_begin:r_end]
        seq_x_mark = self.data_stamp[s_begin:s_end]
        seq_y_mark = self.data_stamp[r_begin:r_end]

        return seq_x, seq_y, seq_x_mark, seq_y_mark

    def __len__(self):
        return len(self.data_x) - self.seq_len - self.pred_len + 1

    def inverse_transform(self, data):
        return self.scaler.inverse_transform(data)
    

class Dataset_Pred(Dataset):
    def __init__(self, root_path, flag='pred', size=None,
                 features='S', data_path='ETTh1.csv',
                 target='OT', scale=True, inverse=False, timeenc=0, freq='15min', cols=None, train_only=False):
        # size [seq_len, label_len, pred_len]
        # info
        if size == None:
            self.seq_len = 24 * 4 * 4
            self.label_len = 24 * 4
            self.pred_len = 24 * 4
        else:
            self.seq_len = size[0]
            self.label_len = size[1]
            self.pred_len = size[2]
        # init
        assert flag in ['pred']

        self.features = features
        self.target = target
        self.scale = scale
        self.inverse = inverse
        self.timeenc = timeenc
        self.freq = freq
        self.cols = cols
        self.root_path = root_path
        self.data_path = data_path
        self.__read_data__()

    def __read_data__(self):
        self.scaler = StandardScaler()
        df_raw = pd.read_csv(os.path.join(self.root_path,
                                          self.data_path))
        '''
        df_raw.columns: ['date', ...(other features), target feature]
        '''
        if self.cols:
            cols = self.cols.copy()
        else:
            cols = list(df_raw.columns)
            self.cols = cols.copy()
            cols.remove('date')
        if self.features == 'S':
            cols.remove(self.target)
        border1 = len(df_raw) - self.seq_len
        border2 = len(df_raw)

        if self.features == 'M' or self.features == 'MS':
            df_raw = df_raw[['date'] + cols]
            cols_data = df_raw.columns[1:]
            df_data = df_raw[cols_data]
        elif self.features == 'S':
            df_raw = df_raw[['date'] + cols + [self.target]]
            df_data = df_raw[[self.target]]

        if self.scale:
            self.scaler.fit(df_data.values)
            data = self.scaler.transform(df_data.values)
        else:
            data = df_data.values

        tmp_stamp = df_raw[['date']][border1:border2]
        tmp_stamp['date'] = pd.to_datetime(tmp_stamp.date)
        pred_dates = pd.date_range(tmp_stamp.date.values[-1], periods=self.pred_len + 1, freq=self.freq)

        df_stamp = pd.DataFrame(columns=['date'])
        df_stamp.date = list(tmp_stamp.date.values) + list(pred_dates[1:])
        self.future_dates = list(pred_dates[1:])
        if self.timeenc == 0:
            df_stamp['month'] = df_stamp.date.apply(lambda row: row.month, 1)
            df_stamp['day'] = df_stamp.date.apply(lambda row: row.day, 1)
            df_stamp['weekday'] = df_stamp.date.apply(lambda row: row.weekday(), 1)
            df_stamp['hour'] = df_stamp.date.apply(lambda row: row.hour, 1)
            df_stamp['minute'] = df_stamp.date.apply(lambda row: row.minute, 1)
            df_stamp['minute'] = df_stamp.minute.map(lambda x: x // 15)
            data_stamp = df_stamp.drop(['date'], 1).values
        elif self.timeenc == 1:
            data_stamp = time_features(pd.to_datetime(df_stamp['date'].values), freq=self.freq)
            data_stamp = data_stamp.transpose(1, 0)

        self.data_x = data[border1:border2]
        if self.inverse:
            self.data_y = df_data.values[border1:border2]
        else:
            self.data_y = data[border1:border2]
        self.data_stamp = data_stamp

    def __getitem__(self, index):
        s_begin = index
        s_end = s_begin + self.seq_len
        r_begin = s_end - self.label_len
        r_end = r_begin + self.label_len + self.pred_len

        seq_x = self.data_x[s_begin:s_end]
        if self.inverse:
            seq_y = self.data_x[r_begin:r_begin + self.label_len]
        else:
            seq_y = self.data_y[r_begin:r_begin + self.label_len]
        seq_x_mark = self.data_stamp[s_begin:s_end]
        seq_y_mark = self.data_stamp[r_begin:r_end]

        return seq_x, seq_y, seq_x_mark, seq_y_mark

    def __len__(self):
        return len(self.data_x) - self.seq_len + 1

    def inverse_transform(self, data):
        return self.scaler.inverse_transform(data)
