from lib.normalization import normalize_dataset

import numpy as np
import gc
import os
import torch
import h5py


def get_dataloader(args, normalizer='std', single=True):
    """STEP模型的数据加载器
    
    Args:
        args: 配置参数
        normalizer: 标准化方法
        single: 是否为单步预测
        
    Returns:
        train_dataloader, val_dataloader, test_dataloader, scaler
    """
    data = load_st_dataset(args['type'], args['sample'])  # 加载数据
    L, N, F = data.shape  # 数据形状

    # Step 1: data -> x,y
    x = add_window_x(data, args['lag'], args['horizon'], single)
    y = add_window_y(data, args['lag'], args['horizon'], single)

    del data
    gc.collect()

    # Step 2: time_in_day, day_in_week -> day, week
    time_in_day = [i % args['steps_per_day'] / args['steps_per_day'] for i in range(L)]
    time_in_day = np.tile(np.array(time_in_day), [1, N, 1]).transpose((2, 1, 0))
    day_in_week = [(i // args['steps_per_day']) % args['days_per_week'] for i in range(L)]
    day_in_week = np.tile(np.array(day_in_week), [1, N, 1]).transpose((2, 1, 0))

    x_day = add_window_x(time_in_day, args['lag'], args['horizon'], single)
    x_week = add_window_x(day_in_week, args['lag'], args['horizon'], single)

    # Step 3 day, week, x, y --> x, y
    x = np.concatenate([x, x_day, x_week], axis=-1)

    del x_day, x_week
    gc.collect()

    # Step 4 x,y --> x_train, x_val, x_test, y_train, y_val, y_test
    if args['test_ratio'] > 1:
        x_train, x_val, x_test = split_data_by_days(x, args['val_ratio'], args['test_ratio'])
    else:
        x_train, x_val, x_test = split_data_by_ratio(x, args['val_ratio'], args['test_ratio'])

    del x
    gc.collect()

    # Normalization
    scaler = normalize_dataset(x_train[..., :args['input_dim']], normalizer, args['column_wise'])
    x_train[..., :args['input_dim']] = scaler.transform(x_train[..., :args['input_dim']])
    x_val[..., :args['input_dim']] = scaler.transform(x_val[..., :args['input_dim']])
    x_test[..., :args['input_dim']] = scaler.transform(x_test[..., :args['input_dim']])

    y_day = add_window_y(time_in_day, args['lag'], args['horizon'], single)
    y_week = add_window_y(day_in_week, args['lag'], args['horizon'], single)

    del time_in_day, day_in_week
    gc.collect()

    y = np.concatenate([y, y_day, y_week], axis=-1)

    del y_day, y_week
    gc.collect()

    # Split Y
    if args['test_ratio'] > 1:
        y_train, y_val, y_test = split_data_by_days(y, args['val_ratio'], args['test_ratio'])
    else:
        y_train, y_val, y_test = split_data_by_ratio(y, args['val_ratio'], args['test_ratio'])

    del y
    gc.collect()

    # Step 5: x_train y_train x_val y_val x_test y_test --> train val test
    train_dataloader = data_loader(x_train, y_train, args['batch_size'], shuffle=True, drop_last=True)

    del x_train, y_train
    gc.collect()

    val_dataloader = data_loader(x_val, y_val, args['batch_size'], shuffle=False, drop_last=True)

    del x_val, y_val
    gc.collect()

    test_dataloader = data_loader(x_test, y_test, args['batch_size'], shuffle=False, drop_last=False)

    del x_test, y_test
    gc.collect()

    return train_dataloader, val_dataloader, test_dataloader, scaler


def load_st_dataset(dataset, sample):
    # output L, N, F
    match dataset:
        case 'PEMSD3':
            data_path = os.path.join('./data/PEMS03/PEMS03.npz')
            data = np.load(data_path)['data']  # (L, N, F)
        case 'PEMSD4':
            data_path = os.path.join('./data/PEMS04/PEMS04.npz')
            data = np.load(data_path)['data']  # (L, N, F)
        case 'PEMSD7':
            data_path = os.path.join('./data/PEMS07/PEMS07.npz')
            data = np.load(data_path)['data']  # (L, N, F)
        case 'PEMSD8':
            data_path = os.path.join('./data/PEMS08/PEMS08.npz')
            data = np.load(data_path)['data']  # (L, N, F)
        case 'METR-LA':
            data_path = os.path.join('./data/METR-LA/METR-LA.npz')
            data = np.load(data_path)['data']  # (L, N, F)
        case 'METR-BAY':
            data_path = os.path.join('./data/METR-BAY/METR-BAY.npz')
            data = np.load(data_path)['data']  # (L, N, F)
        case _:
            raise ValueError(f"Unknown dataset: {dataset}")
    
    if sample:
        data = data[:sample]
    
    return data


def add_window_x(data, lag, horizon, single):
    """
    Add window to data for x
    """
    L, N, F = data.shape
    if single:
        x = np.zeros((L - lag - horizon + 1, lag, N, F))
        for i in range(L - lag - horizon + 1):
            x[i] = data[i:i + lag]
    else:
        x = np.zeros((L - lag - horizon + 1, lag, N, F))
        for i in range(L - lag - horizon + 1):
            x[i] = data[i:i + lag]
    return x


def add_window_y(data, lag, horizon, single):
    """
    Add window to data for y
    """
    L, N, F = data.shape
    if single:
        y = np.zeros((L - lag - horizon + 1, horizon, N, F))
        for i in range(L - lag - horizon + 1):
            y[i] = data[i + lag:i + lag + horizon]
    else:
        y = np.zeros((L - lag - horizon + 1, horizon, N, F))
        for i in range(L - lag - horizon + 1):
            y[i] = data[i + lag:i + lag + horizon]
    return y


def split_data_by_ratio(data, val_ratio, test_ratio):
    """
    Split data by ratio
    """
    L = data.shape[0]
    val_len = int(L * val_ratio)
    test_len = int(L * test_ratio)
    train_len = L - val_len - test_len
    
    train_data = data[:train_len]
    val_data = data[train_len:train_len + val_len]
    test_data = data[train_len + val_len:]
    
    return train_data, val_data, test_data


def split_data_by_days(data, val_days, test_days):
    """
    Split data by days
    """
    L = data.shape[0]
    val_len = val_days * 288  # 288 time steps per day
    test_len = test_days * 288
    train_len = L - val_len - test_len
    
    train_data = data[:train_len]
    val_data = data[train_len:train_len + val_len]
    test_data = data[train_len + val_len:]
    
    return train_data, val_data, test_data


def data_loader(x, y, batch_size, shuffle=True, drop_last=True):
    """
    Create data loader
    """
    dataset = torch.utils.data.TensorDataset(torch.FloatTensor(x), torch.FloatTensor(y))
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last)
    return dataloader