datautils.py

import numpy as np
import pandas as pd
import torch
import json
from torch.utils.data import DataLoader

from src.data.datamodule import DataLoaders
from src.data.pred_dataset import *

# Supported dataset identifiers
SUPPORTED_DATASETS = [
    "PEMS03", "PEMS04", "PEMS07", "PEMS08", 
    "PEMS-BAY", "METR-LA", "Electricity", "Weather", "Traffic", 
    "ETTh1", "ETTh2", "ETTm1", "ETTm2"
]

def get_dls(params):
    """
    Create data loaders with appropriate preprocessing for the specified dataset.
    
    Args:
        params: Configuration object containing:
            dset (str): Dataset identifier from DSETS
            context_points (int): Number of input timesteps
            target_points (int): Number of prediction timesteps 
            features (str): Feature selection mode ('M'=multivariate)
            batch_size (int): Training batch size
            num_workers (int): Number of data loading workers
            finetune_percentage (float): Fraction of data used for fine-tuning
            use_time_features (bool): Whether to include time-based features
    
    Returns:
        DataLoaders: Object containing:
            - train_dl: Training data loader
            - valid_dl: Validation data loader  
            - test_dl: Test data loader
            - vars: Number of variables/features
            - len: Input sequence length
            - c: Output dimension
    
    Notes:
        - All datasets are normalized by default (scale=True)
        - Each dataset uses its own statistics for normalization
        - Data is split into train/val/test sets based on dataset-specific logic
        - Time-based features can be optionally included
    """

    if not hasattr(params, "use_time_features"):
        params.use_time_features = False
    
    if params.dset in SUPPORTED_DATASETS:
        dataset_path = f"data/{params.dset}"
        with open(f'{dataset_path}/desc.json', 'r', encoding='utf-8') as file:
            description = json.load(file)
        size = [params.context_points, 0, params.target_points]
        dls = DataLoaders(
                datasetCls=Dataset_Forall,
                dataset_kwargs={
                    "data_path": f'{dataset_path}/',
                    "data_name": params.dset,
                    "features": params.features,
                    "scale": True,
                    "size": size,
                    "use_time_features": params.use_time_features,
                    "exogenous_variates": getattr(params, 'exogenous_variates', ''),
                    "node_num": description.get('node_num', 1),  # Default to 1 if not specified
                    "past_covariates": description.get('past_covariates', None),
                    "future_covariates": description.get('future_covariates', None),
                    "static_covariates": description.get('static_covariates', None),
                    "finetune_percentage": params.finetune_percentage,
                    "freq": description.get('frequency', 'day'),
                    "split_type": description.get('split_type', 'normal'),
                    "dataset_datesplit": description.get('dataset_datesplit', None),
                    "cycle": getattr(params, 'cycle', 24),
                    "train_size": getattr(params, 'train_size', 'full'),
                },
                batch_size=params.batch_size,
                workers=params.num_workers,
            )
    else:
        raise ValueError(f"Unsupported dataset: {params.dset}. Supported datasets: {SUPPORTED_DATASETS}")
    
    # Store dataset dimensions as attributes
    dls.vars, dls.len = dls.train.dataset[0][0].shape[1], params.context_points
    dls.c = dls.train.dataset[0][1].shape[0]
    return dls


class STDataLoaders(DataLoaders):
    """
    DataLoaders wrapper that augments each dataset split with extra temporal features.

    This keeps the original value channel and stacks additional (categorical) time features
    along a new last dimension.
    """

    def __init__(self, base_dls: DataLoaders, args):
        self.base_dls = base_dls

        dset = getattr(args, "dset", None)
        has_holiday_features = False
        if dset:
            try:
                with open(f"data/{dset}/desc.json", 'r', encoding='utf-8') as f:
                    desc = json.load(f)
                    # Check if holiday_features exists and is not empty
                    if desc.get("holiday_features"):
                        has_holiday_features = True
            except Exception as e:
                print(f"Warning: Could not check holiday features for {dset}: {e}")

        if has_holiday_features:
            train_dataset = add_temporal_features(base_dls.train.dataset)
            valid_dataset = add_temporal_features(base_dls.valid.dataset)
            test_dataset = add_temporal_features(base_dls.test.dataset)
        else:
            train_dataset = add_temporal_features_wo_holiday(base_dls.train.dataset)
            valid_dataset = add_temporal_features_wo_holiday(base_dls.valid.dataset)
            test_dataset = add_temporal_features_wo_holiday(base_dls.test.dataset)

        self.train = DataLoader(
            train_dataset,
            batch_size=args.batch_size,
            shuffle=False,
            num_workers=args.num_workers,
            drop_last=True,
        )

        self.valid = DataLoader(
            valid_dataset,
            batch_size=args.batch_size,
            shuffle=False,
            num_workers=args.num_workers,
            drop_last=True,
        )

        self.test = DataLoader(
            test_dataset,
            batch_size=args.batch_size,
            shuffle=False,
            num_workers=args.num_workers,
            drop_last=False,
        )

        # Basic dimension metadata (kept for compatibility with training scripts)
        self.vars = train_dataset.data_x.shape[1]  # number of nodes / variables
        self.len = len(self.train)
        self.c = train_dataset.data_y.shape[-1] if len(train_dataset.data_y.shape) > 2 else 1


def get_st_dls(args):
    """
    Build dataloaders and augment them with spatio-temporal (time) features.
    """
    base_dls = get_dls(args)
    original_x_shape = base_dls.train.dataset.data_x.shape
    original_y_shape = base_dls.train.dataset.data_y.shape

    dls = STDataLoaders(base_dls, args)

    print(
        f"Added temporal features to train data_x shape: {original_x_shape} → {dls.train.dataset.data_x.shape}"
    )
    print(
        f"Added temporal features to train data_y shape: {original_y_shape} → {dls.train.dataset.data_y.shape}"
    )
    return dls


def _ensure_numpy(x):
    if isinstance(x, torch.Tensor):
        return x.detach().cpu().numpy()
    return x


def add_temporal_features_wo_holiday(dataset):
    """
    Add basic timestamp-derived features (minute/hour/weekday/day/month) to dataset.data_x.
    """
    data_x = dataset.data_x  # [samples, nodes]
    data_y = dataset.data_y  # [samples, nodes]

    _, num_nodes = data_x.shape

    # Note: Dataset_Forall uses [month, day, weekday, hour, minutes, ...] when timeenc==0.
    minutes_of_hour = dataset.data_stamp[:, 4]
    time_of_day = dataset.data_stamp[:, 3]
    day_of_week = dataset.data_stamp[:, 2]
    day_of_month = dataset.data_stamp[:, 1]
    month_of_year = dataset.data_stamp[:, 0]

    minutes_of_hour = minutes_of_hour.reshape(-1, 1).repeat(num_nodes, axis=1)
    time_of_day = time_of_day.reshape(-1, 1).repeat(num_nodes, axis=1)
    day_of_week = day_of_week.reshape(-1, 1).repeat(num_nodes, axis=1)
    day_of_month = day_of_month.reshape(-1, 1).repeat(num_nodes, axis=1)
    month_of_year = month_of_year.reshape(-1, 1).repeat(num_nodes, axis=1)

    data_x_np = _ensure_numpy(data_x)

    data_x_new = np.stack(
        [
            data_x_np,
            minutes_of_hour,
            time_of_day,
            day_of_week,
            day_of_month,
            month_of_year,
        ],
        axis=2,
    )  # [samples, nodes, 6]

    if isinstance(dataset.data_x, torch.Tensor):
        data_x_tensor = torch.from_numpy(data_x_new).float()
        data_x_tensor[..., 1:] = data_x_tensor[..., 1:].long()
        if dataset.data_x.is_cuda:
            data_x_tensor = data_x_tensor.cuda()
        data_x_new = data_x_tensor

    dataset.data_x = data_x_new
    dataset.data_y = data_y
    return dataset


def add_temporal_features(dataset):
    """
    Add basic timestamp-derived features + holiday-related categorical features to dataset.data_x.
    """
    data_x = dataset.data_x  # [samples, nodes]
    data_y = dataset.data_y  # [samples, nodes]

    _, num_nodes = data_x.shape

    # Dataset_Forall (timeenc==0) produces:
    # [month, day, weekday, hour, minutes, spring_festival, task_holiday, statutory_date, name]
    name = dataset.data_stamp[:, 8]
    statutory_date = dataset.data_stamp[:, 7]
    task_holiday = dataset.data_stamp[:, 6]
    spring_festival = dataset.data_stamp[:, 5]
    minutes_of_hour = dataset.data_stamp[:, 4]
    time_of_day = dataset.data_stamp[:, 3]
    day_of_week = dataset.data_stamp[:, 2]
    day_of_month = dataset.data_stamp[:, 1]
    month_of_year = dataset.data_stamp[:, 0]

    name = name.reshape(-1, 1).repeat(num_nodes, axis=1)
    statutory_date = statutory_date.reshape(-1, 1).repeat(num_nodes, axis=1)
    task_holiday = task_holiday.reshape(-1, 1).repeat(num_nodes, axis=1)
    spring_festival = spring_festival.reshape(-1, 1).repeat(num_nodes, axis=1)
    minutes_of_hour = minutes_of_hour.reshape(-1, 1).repeat(num_nodes, axis=1)
    time_of_day = time_of_day.reshape(-1, 1).repeat(num_nodes, axis=1)
    day_of_week = day_of_week.reshape(-1, 1).repeat(num_nodes, axis=1)
    day_of_month = day_of_month.reshape(-1, 1).repeat(num_nodes, axis=1)
    month_of_year = month_of_year.reshape(-1, 1).repeat(num_nodes, axis=1)

    data_x_np = _ensure_numpy(data_x)

    data_x_new = np.stack(
        [
            data_x_np,
            minutes_of_hour,
            time_of_day,
            day_of_week,
            day_of_month,
            month_of_year,
            spring_festival,
            task_holiday,
            statutory_date,
            name,
        ],
        axis=2,
    )  # [samples, nodes, 10]

    if isinstance(dataset.data_x, torch.Tensor):
        data_x_tensor = torch.from_numpy(data_x_new).float()
        data_x_tensor[..., 1:] = data_x_tensor[..., 1:].long()
        if dataset.data_x.is_cuda:
            data_x_tensor = data_x_tensor.cuda()
        data_x_new = data_x_tensor

    dataset.data_x = data_x_new
    dataset.data_y = data_y
    return dataset