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):
    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[..., :args['input_dim']], y_train[..., :args['input_dim']], args['batch_size'], shuffle=True, drop_last=True)
    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[..., :args['input_dim']], y_val[..., :args['input_dim']], args['batch_size'], shuffle=False, drop_last=True)
    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[..., :args['input_dim']], y_test[..., :args['input_dim']], args['batch_size'], shuffle=False, drop_last=False)
    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 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 _:
            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)