import numpy as np
import gc
import os
import torch
import h5py
from lib.normalization import normalize_dataset


def get_dataloader(args, normalizer='std', single=True):
    # args should now include 'cycle'
    data = load_st_dataset(args['type'], args['sample'])  # [T, N, F]
    L, N, F = data.shape

    # compute cycle index
    cycle_arr = np.arange(L) % args['cycle']  # length-L array

    # Step 1: sliding windows for X and Y
    x = add_window_x(data, args['lag'], args['horizon'], single)
    y = add_window_y(data, args['lag'], args['horizon'], single)
    # window count = M = L - lag - horizon + 1
    M = x.shape[0]

    # Step 2: time features
    time_in_day = np.tile(
        np.array([i % args['steps_per_day'] / args['steps_per_day'] for i in range(L)]),
        (N, 1)
    ).T.reshape(L, N, 1)
    day_in_week = np.tile(
        np.array([(i // args['steps_per_day']) % args['days_per_week'] for i in range(L)]),
        (N, 1)
    ).T.reshape(L, N, 1)

    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)
    x = np.concatenate([x, x_day, x_week], axis=-1)
    # del x_day, x_week
    # gc.collect()

    # Step 3: extract cycle index per window: take value at end of sequence
    cycle_win = np.array([cycle_arr[i + args['lag']] for i in range(M)])  # shape [M]

    # Step 4: split into train/val/test
    if args['test_ratio'] > 1:
        x_train, x_val, x_test = split_data_by_days(x, args['val_ratio'], args['test_ratio'])
        y_train, y_val, y_test = split_data_by_days(y, args['val_ratio'], args['test_ratio'])
        c_train, c_val, c_test = split_data_by_days(cycle_win, args['val_ratio'], args['test_ratio'])
    else:
        x_train, x_val, x_test = split_data_by_ratio(x, args['val_ratio'], args['test_ratio'])
        y_train, y_val, y_test = split_data_by_ratio(y, args['val_ratio'], args['test_ratio'])
        c_train, c_val, c_test = split_data_by_ratio(cycle_win, args['val_ratio'], args['test_ratio'])
    # del x, y, cycle_win
    # gc.collect()

    # Step 5: normalization on X only
    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']])

    # add time features to Y
    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)
    y = np.concatenate([y, y_day, y_week], axis=-1)
    # del y_day, y_week, time_in_day, day_in_week
    # gc.collect()

    # split Y time-augmented
    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

    # Step 6: create dataloaders including cycle index
    train_loader = data_loader_with_cycle(x_train, y_train, c_train, args['batch_size'], shuffle=True, drop_last=True)
    val_loader = data_loader_with_cycle(x_val, y_val, c_val, args['batch_size'], shuffle=False, drop_last=True)
    test_loader = data_loader_with_cycle(x_test, y_test, c_test, args['batch_size'], shuffle=False, drop_last=False)

    return train_loader, val_loader, test_loader, scaler


def data_loader_with_cycle(X, Y, C, batch_size, shuffle=True, drop_last=True):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    X_t = torch.tensor(X, dtype=torch.float32, device=device)
    Y_t = torch.tensor(Y, dtype=torch.float32, device=device)
    C_t = torch.tensor(C, dtype=torch.long, device=device).unsqueeze(-1)  # [B,1]
    dataset = torch.utils.data.TensorDataset(X_t, Y_t, C_t)
    loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last)
    return loader

# Rest of the helper functions (load_st_dataset, split_data..., add_window_x/y) unchanged


def load_st_dataset(dataset, sample):
    # output B, N, D
    match dataset:
        case 'PEMSD3':
            data_path = os.path.join('./data/PEMS03/PEMS03.npz')
            data = np.load(data_path)['data'][:, :, 0]  # only the first dimension, traffic flow data
        case 'PEMSD4':
            data_path = os.path.join('./data/PEMS04/PEMS04.npz')
            data = np.load(data_path)['data'][:, :, 0]  # only the first dimension, traffic flow data
        case 'PEMSD7':
            data_path = os.path.join('./data/PEMS07/PEMS07.npz')
            data = np.load(data_path)['data'][:, :, 0]  # only the first dimension, traffic flow data
        case 'PEMSD8':
            data_path = os.path.join('./data/PEMS08/PEMS08.npz')
            data = np.load(data_path)['data'][:, :, 0]  # only the first dimension, traffic flow data
        case 'PEMSD7(L)':
            data_path = os.path.join('./data/PEMS07(L)/PEMS07L.npz')
            data = np.load(data_path)['data'][:, :, 0]  # only the first dimension, traffic flow data
        case 'PEMSD7(M)':
            data_path = os.path.join('./data/PEMS07(M)/V_228.csv')
            data = np.genfromtxt(data_path, delimiter=',')  # Read CSV directly with numpy
        case 'METR-LA':
            data_path = os.path.join('./data/METR-LA/METR.h5')
            with h5py.File(data_path, 'r') as f:  # Use h5py to handle HDF5 files without pandas
                data = np.array(f['data'])
        case 'BJ':
            data_path = os.path.join('./data/BJ/BJ500.csv')
            data = np.genfromtxt(data_path, delimiter=',', skip_header=1)  # Skip header if present
        case 'Hainan':
            data_path = os.path.join('./data/Hainan/Hainan.npz')
            data = np.load(data_path)['data'][:, :, 0]
        case 'SD':
            data_path = os.path.join('./data/SD/data.npz')
            data = np.load(data_path)["data"][:, :, 0].astype(np.float32)
        case _:
            raise ValueError(f"Unsupported dataset: {dataset}")

    # Ensure data shape compatibility
    if len(data.shape) == 2:
        data = np.expand_dims(data, axis=-1)

    print('加载 %s 数据集中... ' % dataset)
    return data[::sample]

def split_data_by_days(data, val_days, test_days, interval=30):
    t = int((24 * 60) / interval)
    test_data = data[-t * int(test_days):]
    val_data = data[-t * int(test_days + val_days):-t * int(test_days)]
    train_data = data[:-t * int(test_days + val_days)]
    return train_data, val_data, test_data


def split_data_by_ratio(data, val_ratio, test_ratio):
    data_len = data.shape[0]
    test_data = data[-int(data_len * test_ratio):]
    val_data = data[-int(data_len * (test_ratio + val_ratio)):-int(data_len * test_ratio)]
    train_data = data[:-int(data_len * (test_ratio + val_ratio))]
    return train_data, val_data, test_data


def data_loader(X, Y, batch_size, shuffle=True, drop_last=True):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    X = torch.tensor(X, dtype=torch.float32, device=device)
    Y = torch.tensor(Y, dtype=torch.float32, device=device)
    data = torch.utils.data.TensorDataset(X, Y)
    dataloader = torch.utils.data.DataLoader(data, batch_size=batch_size,
                                             shuffle=shuffle, drop_last=drop_last)
    return dataloader


def add_window_x(data, window=3, horizon=1, single=False):
    """
    Generate windowed X values from the input data.

    :param data: Input data, shape [B, ...]
    :param window: Size of the sliding window
    :param horizon: Horizon size
    :param single: If True, generate single-step windows, else multi-step
    :return: X with shape [B, W, ...]
    """
    length = len(data)
    end_index = length - horizon - window + 1
    x = []  # Sliding windows
    index = 0

    while index < end_index:
        x.append(data[index:index + window])
        index += 1

    return np.array(x)


def add_window_y(data, window=3, horizon=1, single=False):
    """
    Generate windowed Y values from the input data.

    :param data: Input data, shape [B, ...]
    :param window: Size of the sliding window
    :param horizon: Horizon size
    :param single: If True, generate single-step windows, else multi-step
    :return: Y with shape [B, H, ...]
    """
    length = len(data)
    end_index = length - horizon - window + 1
    y = []  # Horizon values
    index = 0

    while index < end_index:
        if single:
            y.append(data[index + window + horizon - 1:index + window + horizon])
        else:
            y.append(data[index + window:index + window + horizon])
        index += 1

    return np.array(y)


if __name__ == '__main__':
    res = load_st_dataset('SD', 1)
    k = 1