run_baselines.py

"""
Unified baseline training & evaluation script for traffic speed prediction.

Usage examples:
    # Train LSTM on I4_Hills dataset
    python run_baselines.py --model LSTM --dataset I4_Hills --mode train

    # Test a trained DCRNN model on I75_Hills dataset
    python run_baselines.py --model DCRNN --dataset I75_Hills --mode test

    # Train and test all models on all datasets
    python run_baselines.py --model all --dataset all --mode both

    # Override hyperparameters
    python run_baselines.py --model STGCN --dataset I4_Orange --mode train --epochs 200 --lr 0.001

Supported models: LSTM, DCRNN, STGCN, GraphWaveNet, ASTGCN, iTransformer, SATEformer, TLAST
Supported datasets: I4_Hills, I75_Hills, I275_Hills, I4_Orange
"""

import argparse
import os
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset, random_split
from sklearn.metrics import mean_squared_error, mean_absolute_error

from baselines import (
    Seq2SeqLSTM,
    DCRNN,
    STGCN,
    GraphWaveNet,
    make_astgcn,
    SATEformer,
    TLASTWrapper,
    _tlast_args,
    _import_tlast_model,
)

# ========================= Constants =========================

SPEED_MIN, SPEED_MAX = 0.0, 101.5
CHANNELS_TO_DELETE = [0, 4, 5, 6, 7, 11]
# iTransformer only keeps channel index 1
ITRANSFORMER_CHANNELS_TO_DELETE = [0, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13]
CH_SPEED = 1
CH_HOUR = 12
CH_WEEKDAY = 13

ALL_MODELS = [
    "LSTM", "DCRNN", "STGCN", "GraphWaveNet", "ASTGCN", "iTransformer", "SATEformer", "TLAST",
]

# Mapping from model name to the checkpoint file name suffix
# (to match checkpoint files in checkpoints/baselines/)
CKPT_NAME_MAP = {
    "LSTM": "LSTM",
    "DCRNN": "DCRNN",
    "STGCN": "STGCN",
    "GraphWaveNet": "GraphWave",
    "ASTGCN": "ASTGCN",
    "iTransformer": "iTransformer",
    "SATEformer": "SATEformer",
    "TLAST": "TLAST",
}

DATASET_CONFIG = {
    "I4_Hills": {
        "train": "data/I4_Hillsborough_train_processed.npz",
        "test": "data/I4_Hillsborough_test_processed.npz",
        "adj": "data/Adj_Hills_I4.pt",
        "prefix": "I4_Hills",
    },
    "I75_Hills": {
        "train": "data/I75_Hillsborough_train_processed.npz",
        "test": "data/I75_Hillsborough_test_processed.npz",
        "adj": "data/Adj_Hills_I75.pt",
        "prefix": "I75_Hills",
    },
    "I275_Hills": {
        "train": "data/I275_Hillsborough_train_processed.npz",
        "test": "data/I275_Hillsborough_test_processed.npz",
        "adj": "data/Adj_Hills_I275.pt",
        "prefix": "I275_Hills",
    },
    "I4_Orange": {
        "train": "data/I4_Orange_Seminole_train_processed.npz",
        "test": "data/I4_Orange_Seminole_test_processed.npz",
        "adj": "data/Adj_Orange_I4.pt",
        "prefix": "I4_Orange",
    },
}

MODEL_DEFAULTS = {
    "LSTM":          {"batch_size": 64, "lr": 1e-3, "patience": 5},
    "DCRNN":         {"batch_size": 64, "lr": 1e-2, "patience": 5},
    "STGCN":         {"batch_size": 64, "lr": 5e-3, "patience": 10},
    "GraphWaveNet":  {"batch_size": 64, "lr": 5e-3, "patience": 5},
    "ASTGCN":        {"batch_size": 64, "lr": 5e-3, "patience": 5},
    "iTransformer":  {"batch_size": 32, "lr": 1e-2, "patience": 10},
    "SATEformer":    {"batch_size": 32, "lr": 1e-3, "patience": 5},
    # Official PEMS04 12->12-style TLAST hparams used in the rebuttal runs
    "TLAST":         {"batch_size": 32, "lr": 1e-3, "patience": 20},
}

# SATEformer: only I4_Hills uses tod/dow embeddings
SATEFORMER_TOD_DOW = {
    "I4_Hills":  {"tod_embedding_dim": 2, "dow_embedding_dim": 2},
    "I75_Hills": {"tod_embedding_dim": 0, "dow_embedding_dim": 0},
    "I275_Hills": {"tod_embedding_dim": 0, "dow_embedding_dim": 0},
    "I4_Orange": {"tod_embedding_dim": 0, "dow_embedding_dim": 0},
}


# ========================= Utility =========================

def inverse_transform(y_norm):
    return y_norm * (SPEED_MAX - SPEED_MIN) + SPEED_MIN


def mape(y_true, y_pred, eps=1e-3):
    return np.mean(np.abs((y_true - y_pred) / (y_true + eps))) * 100


def get_checkpoint_path(save_dir, prefix, model_name):
    ckpt_suffix = CKPT_NAME_MAP.get(model_name, model_name)
    return os.path.join(save_dir, f"{prefix}_{ckpt_suffix}.pth")


# ========================= Model Definitions =========================

# ========================= Data Loading =========================

def load_raw_data(dataset_name, data_dir):
    cfg = DATASET_CONFIG[dataset_name]
    train_path = os.path.join(data_dir, cfg["train"]) if data_dir else cfg["train"]
    test_path = os.path.join(data_dir, cfg["test"]) if data_dir else cfg["test"]

    with np.load(train_path) as d:
        X_train, y_train = d["sequences"], d["targets"]
    with np.load(test_path) as d:
        X_test, y_test = d["sequences"], d["targets"]

    return X_train, y_train, X_test, y_test


def load_adj(dataset_name, device, data_dir=None):
    cfg = DATASET_CONFIG[dataset_name]
    adj_path = os.path.join(data_dir, cfg["adj"]) if data_dir else cfg["adj"]
    return torch.load(adj_path, map_location=device, weights_only=True).to(device)


def preprocess_standard(X_train, y_train, X_test, y_test, device):
    """Standard preprocessing: drop channels, min-max normalize all. Returns (B,T,N,C) format."""
    X_tr = torch.tensor(X_train, dtype=torch.float32, device=device)
    y_tr = torch.tensor(y_train, dtype=torch.float32, device=device)
    X_te = torch.tensor(X_test, dtype=torch.float32, device=device)
    y_te = torch.tensor(y_test, dtype=torch.float32, device=device)

    keep_idx = torch.tensor(
        [i for i in range(X_tr.size(-1)) if i not in CHANNELS_TO_DELETE], device=device)
    X_tr = X_tr.index_select(-1, keep_idx)
    X_te = X_te.index_select(-1, keep_idx)

    min_vals = X_tr.amin(dim=(0, 1, 2), keepdim=True)
    max_vals = X_tr.amax(dim=(0, 1, 2), keepdim=True)
    rng = (max_vals - min_vals).clamp_min(1e-6)
    X_tr = (X_tr - min_vals) / rng
    X_te = (X_te - min_vals) / rng

    return X_tr, y_tr, X_te, y_te


def preprocess_tlast(X_train, y_train, X_test, y_test, device):
    """TLAST: z-score speed (mph) + shared hour/weekday; targets kept in mph."""
    def _calendar(X):
        # Match rebuttal adapter: calendar from node 0, then broadcast to N
        hour = np.rint(X[:, :, 0, CH_HOUR]).astype(np.float32) % 24.0
        weekday = np.rint(X[:, :, 0, CH_WEEKDAY]).astype(np.float32)
        weekday = np.where(weekday >= 1.0, weekday - 1.0, weekday)
        weekday = np.clip(weekday, 0.0, 6.0)
        return hour, weekday

    speed_tr = inverse_transform(X_train[..., CH_SPEED]).astype(np.float32)
    speed_te = inverse_transform(X_test[..., CH_SPEED]).astype(np.float32)
    mean = float(speed_tr.mean())
    std = float(speed_tr.std())
    if std < 1e-6:
        std = 1.0

    def _x(speed, X):
        hour, weekday = _calendar(X)
        hour = np.broadcast_to(hour[:, :, None], speed.shape)
        weekday = np.broadcast_to(weekday[:, :, None], speed.shape)
        speed_z = (speed - mean) / std
        return np.stack([speed_z, hour, weekday], axis=-1).astype(np.float32)

    y_tr = inverse_transform(y_train).astype(np.float32)
    y_te = inverse_transform(y_test).astype(np.float32)
    X_tr = torch.tensor(_x(speed_tr, X_train), dtype=torch.float32, device=device)
    X_te = torch.tensor(_x(speed_te, X_test), dtype=torch.float32, device=device)
    y_tr = torch.tensor(y_tr, dtype=torch.float32, device=device)
    y_te = torch.tensor(y_te, dtype=torch.float32, device=device)
    return X_tr, y_tr, X_te, y_te, mean, std

def preprocess_sateformer(X_train, y_train, X_test, y_test, device):
    """SATEformer preprocessing: separate continuous/discrete, normalize only continuous."""
    X_tr = torch.tensor(X_train, dtype=torch.float32, device=device)
    y_tr = torch.tensor(y_train, dtype=torch.float32, device=device)
    X_te = torch.tensor(X_test, dtype=torch.float32, device=device)
    y_te = torch.tensor(y_test, dtype=torch.float32, device=device)

    keep_idx = torch.tensor(
        [i for i in range(X_tr.size(-1)) if i not in CHANNELS_TO_DELETE], device=device)
    X_tr = X_tr.index_select(-1, keep_idx)
    X_te = X_te.index_select(-1, keep_idx)

    X_tr_cont, X_tr_disc = X_tr[..., :-2], X_tr[..., -2:]
    X_te_cont, X_te_disc = X_te[..., :-2], X_te[..., -2:]

    X_tr_disc[..., 0] = X_tr_disc[..., 0].long()
    X_tr_disc[..., 1] = (X_tr_disc[..., 1] - 1).clamp(min=0, max=6).long()
    X_te_disc[..., 0] = X_te_disc[..., 0].long()
    X_te_disc[..., 1] = (X_te_disc[..., 1] - 1).clamp(min=0, max=6).long()

    min_vals = X_tr_cont.amin(dim=(0, 1, 2), keepdim=True)
    max_vals = X_tr_cont.amax(dim=(0, 1, 2), keepdim=True)
    rng = (max_vals - min_vals).clamp_min(1e-6)
    X_tr_cont = (X_tr_cont - min_vals) / rng
    X_te_cont = (X_te_cont - min_vals) / rng

    X_tr = torch.cat([X_tr_cont, X_tr_disc], dim=-1)
    X_te = torch.cat([X_te_cont, X_te_disc], dim=-1)

    y_tr = y_tr.unsqueeze(-1)
    y_te = y_te.unsqueeze(-1)

    return X_tr, y_tr, X_te, y_te


def preprocess_itransformer(X_train, y_train, X_test, y_test, device):
    """iTransformer preprocessing: only keep channel 1, flatten (N*C)."""
    X_tr = torch.tensor(X_train, dtype=torch.float32, device=device)
    y_tr = torch.tensor(y_train, dtype=torch.float32, device=device)
    X_te = torch.tensor(X_test, dtype=torch.float32, device=device)
    y_te = torch.tensor(y_test, dtype=torch.float32, device=device)

    keep_idx = torch.tensor(
        [i for i in range(X_tr.size(-1)) if i not in ITRANSFORMER_CHANNELS_TO_DELETE],
        device=device)
    X_tr = X_tr.index_select(-1, keep_idx)
    X_te = X_te.index_select(-1, keep_idx)

    min_vals = X_tr.amin(dim=(0, 1, 2), keepdim=True)
    max_vals = X_tr.amax(dim=(0, 1, 2), keepdim=True)
    rng = (max_vals - min_vals).clamp_min(1e-6)
    X_tr = (X_tr - min_vals) / rng
    X_te = (X_te - min_vals) / rng

    B, T, N, C = X_tr.shape
    X_tr = X_tr.reshape(B, T, N * C)
    X_te = X_te.reshape(X_te.shape[0], T, N * C)
    y_tr = y_tr.reshape(y_tr.shape[0], T, -1)
    y_te = y_te.reshape(y_te.shape[0], T, -1)

    return X_tr, y_tr, X_te, y_te, N, C


# ========================= Model Construction =========================

def build_model(model_name, dataset_name, device, data_shapes, adj=None):
    """Build and return the model given name, shapes, and optional adjacency matrix.
    
    data_shapes: dict with keys like 'N', 'C', 'T', 'H' (horizon), etc.
    """
    N = data_shapes["N"]
    C = data_shapes["C"]
    T = data_shapes["T"]
    H = data_shapes["H"]

    if model_name == "LSTM":
        return Seq2SeqLSTM(num_features=C, num_segments=N, hidden_dim=128).to(device)

    elif model_name == "DCRNN":
        return DCRNN(N, C, hid=32, K=2, adj=adj, horizon=H).to(device)

    elif model_name == "STGCN":
        return STGCN(in_channels=C, num_nodes=N, horizon=H, adj=adj).to(device)

    elif model_name == "GraphWaveNet":
        return GraphWaveNet(
            device=device, num_nodes=N, in_dim=C, out_dim=H, supports=[adj]
        ).to(device)

    elif model_name == "ASTGCN":
        return make_astgcn(
            device=device, adj_mx=adj, in_channels=C,
            num_vertices=N, len_input=T
        )

    elif model_name == "iTransformer":
        from iTransformer import iTransformer as iTransformerModel
        return iTransformerModel(
            num_variates=data_shapes["N_times_C"],
            lookback_len=T,
            dim=64, depth=3, heads=4, dim_head=64,
            pred_length=(T,),
            num_tokens_per_variate=1,
            use_reversible_instance_norm=True
        ).to(device)

    elif model_name == "SATEformer":
        tod_dow = SATEFORMER_TOD_DOW.get(dataset_name, {"tod_embedding_dim": 0, "dow_embedding_dim": 0})
        return SATEformer(
            num_nodes=N, in_steps=T, out_steps=T,
            input_dim=C - 2, output_dim=1, steps_per_day=24,
            input_embedding_dim=24,
            tod_embedding_dim=tod_dow["tod_embedding_dim"],
            dow_embedding_dim=tod_dow["dow_embedding_dim"],
            spatial_embedding_dim=0, adaptive_embedding_dim=40,
            feed_forward_dim=128, num_heads=4, num_layers=3,
            dropout=0.1, use_mixed_proj=True
        ).to(device)

    elif model_name == "TLAST":
        TLASTModel = _import_tlast_model()
        core = TLASTModel(_tlast_args(N))
        for p in core.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)
        return TLASTWrapper(
            core,
            scaler_mean=data_shapes["scaler_mean"],
            scaler_std=data_shapes["scaler_std"],
        ).to(device)

    else:
        raise ValueError(f"Unknown model: {model_name}")


# ========================= Training =========================

def train_model(model, model_name, train_loader, val_loader, device,
                checkpoint_path, data_shapes, epochs=100, lr=None, patience=None):
    defaults = MODEL_DEFAULTS[model_name]
    lr = lr or defaults["lr"]
    patience = patience or defaults["patience"]

    if model_name == "TLAST":
        criterion = nn.HuberLoss(delta=2.0)
        optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4)
        scheduler = optim.lr_scheduler.CosineAnnealingLR(
            optimizer, T_max=epochs, eta_min=0.0)
    else:
        criterion = nn.MSELoss()
        optimizer = optim.Adam(model.parameters(), lr=lr,
                               weight_decay=1e-3 if model_name == "iTransformer" else 0)
        scheduler = None
        if model_name == "iTransformer":
            from torch.optim.lr_scheduler import ReduceLROnPlateau
            scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=5)

    best_mae = float('inf')
    wait = 0
    H = data_shapes["H"]
    T = data_shapes["T"]
    N_orig = data_shapes.get("N_orig", data_shapes["N"])
    already_mph = model_name == "TLAST"

    for epoch in range(1, epochs + 1):
        model.train()
        train_loss = 0.0

        for xb, yb in train_loader:
            optimizer.zero_grad(set_to_none=already_mph)

            if model_name == "LSTM":
                out = model(xb, target_len=H)
            elif model_name == "iTransformer":
                pred_dict = model(xb)
                out = pred_dict[T][:, :, :N_orig]
            else:
                out = model(xb)

            loss = criterion(out, yb)
            loss.backward()
            if already_mph:
                nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            train_loss += loss.item()

        train_loss /= len(train_loader)
        if model_name == "TLAST":
            scheduler.step()

        # Validation
        model.eval()
        preds, trues = [], []
        with torch.no_grad():
            for xb, yb in val_loader:
                if model_name == "LSTM":
                    out = model(xb, target_len=H)
                elif model_name == "iTransformer":
                    out = model(xb)[T][:, :, :N_orig]
                else:
                    out = model(xb)
                preds.append(out.cpu().numpy())
                trues.append(yb.cpu().numpy())

        y_val = np.concatenate(trues, 0)
        p_val = np.concatenate(preds, 0)
        if not already_mph:
            y_val = inverse_transform(y_val)
            p_val = inverse_transform(p_val)
        mae_val = mean_absolute_error(y_val.flatten(), p_val.flatten())

        if scheduler is not None and model_name != "TLAST":
            scheduler.step(mae_val)

        if mae_val < best_mae:
            best_mae = mae_val
            wait = 0
            os.makedirs(os.path.dirname(checkpoint_path) or ".", exist_ok=True)
            # checkpoints/baselines TLAST ckpts store the inner Model state_dict only
            state = model.core.state_dict() if model_name == "TLAST" else model.state_dict()
            torch.save(state, checkpoint_path)
        else:
            wait += 1
            if wait >= patience:
                print(f"  Early stopping at epoch {epoch} (no improvement for {patience} epochs).")
                break

        lr_info = f", LR={optimizer.param_groups[0]['lr']:.6f}" if scheduler else ""
        print(f"  Epoch {epoch}/{epochs}: TrainLoss={train_loss:.4f}, ValMAE={mae_val:.4f}{lr_info}")

    print(f"  Best ValMAE: {best_mae:.4f}")
    return best_mae


# ========================= Evaluation =========================

def evaluate_model(model, model_name, X_test_tensor, y_test_np, device,
                   checkpoint_path, data_shapes):
    H = data_shapes["H"]
    T = data_shapes["T"]
    N_orig = data_shapes.get("N_orig", data_shapes["N"])

    state = torch.load(checkpoint_path, map_location=device, weights_only=True)
    if model_name == "TLAST":
        model.core.load_state_dict(state)
    else:
        model.load_state_dict(state)
    model.to(device)
    model.eval()

    with torch.no_grad():
        x = X_test_tensor.to(device)
        if model_name == "LSTM":
            pred_test = model(x, target_len=H).cpu().numpy()
        elif model_name == "iTransformer":
            pred_test = model(x)[T][:, :, :N_orig].cpu().numpy()
        else:
            pred_test = model(x).cpu().numpy()

    # ASTGCN output is (B, N, T) — transpose to (B, T, N) for unified evaluation
    if model_name == "ASTGCN":
        pred_test = pred_test.transpose(0, 2, 1)
        y_test_np = y_test_np.transpose(0, 2, 1)

    # TLAST targets / preds are already in mph
    if model_name == "TLAST":
        y_test_orig = y_test_np
        pred_test_orig = pred_test
    else:
        y_test_orig = inverse_transform(y_test_np)
        pred_test_orig = inverse_transform(pred_test)

    rmse = np.sqrt(mean_squared_error(y_test_orig.flatten(), pred_test_orig.flatten()))
    mae_val = mean_absolute_error(y_test_orig.flatten(), pred_test_orig.flatten())
    mape_val = mape(y_test_orig.flatten(), pred_test_orig.flatten())

    print(f"  Overall => RMSE: {rmse:.4f}, MAE: {mae_val:.4f}, MAPE: {mape_val:.4f}")

    n_steps = y_test_orig.shape[1] if y_test_orig.ndim >= 2 else 1
    if n_steps > 1:
        print("  Per-step performance:")
        for step in range(n_steps):
            if y_test_orig.ndim == 3:
                y_s = y_test_orig[:, step, :].flatten()
                p_s = pred_test_orig[:, step, :].flatten()
            elif y_test_orig.ndim == 4:
                y_s = y_test_orig[:, step, :, :].flatten()
                p_s = pred_test_orig[:, step, :, :].flatten()
            else:
                break
            rmse_s = np.sqrt(mean_squared_error(y_s, p_s))
            mae_s = mean_absolute_error(y_s, p_s)
            mape_s = mape(y_s, p_s)
            print(f"    Step {step + 1:2d}: RMSE={rmse_s:.4f}, MAE={mae_s:.4f}, MAPE={mape_s:.2f}%")

    return {"rmse": rmse, "mae": mae_val, "mape": mape_val}


# ========================= Main Pipeline =========================

def run_single(model_name, dataset_name, mode, args):
    print(f"\n{'=' * 60}")
    print(f"  Model: {model_name} | Dataset: {dataset_name} | Mode: {mode}")
    print(f"{'=' * 60}")

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"  Device: {device}")

    prefix = DATASET_CONFIG[dataset_name]["prefix"]
    train_ckpt_path = get_checkpoint_path(args.save_dir, prefix, model_name)
    test_ckpt_path = get_checkpoint_path(args.load_dir, prefix, model_name)

    # Load raw data
    X_train, y_train, X_test, y_test = load_raw_data(dataset_name, args.data_dir)

    # Load adjacency if needed
    needs_adj = model_name in ("DCRNN", "STGCN", "GraphWaveNet", "ASTGCN")
    adj = load_adj(dataset_name, device, args.data_dir) if needs_adj else None

    # Preprocess based on model
    if model_name == "iTransformer":
        X_tr, y_tr, X_te, y_te, N_orig, C_orig = preprocess_itransformer(
            X_train, y_train, X_test, y_test, device)
        T = X_tr.shape[1]
        data_shapes = {
            "N": X_tr.shape[-1], "C": 1, "T": T,
            "H": y_tr.shape[1], "N_times_C": X_tr.shape[-1], "N_orig": N_orig
        }
    elif model_name == "SATEformer":
        X_tr, y_tr, X_te, y_te = preprocess_sateformer(
            X_train, y_train, X_test, y_test, device)
        T, N, C = X_tr.shape[1], X_tr.shape[2], X_tr.shape[3]
        data_shapes = {"N": N, "C": C, "T": T, "H": T}
    elif model_name == "TLAST":
        X_tr, y_tr, X_te, y_te, scaler_mean, scaler_std = preprocess_tlast(
            X_train, y_train, X_test, y_test, device)
        T, N, C = X_tr.shape[1], X_tr.shape[2], X_tr.shape[3]
        data_shapes = {
            "N": N, "C": C, "T": T, "H": y_tr.shape[1],
            "scaler_mean": scaler_mean, "scaler_std": scaler_std,
        }
    else:
        X_tr, y_tr, X_te, y_te = preprocess_standard(
            X_train, y_train, X_test, y_test, device)

        N, C = X_tr.shape[2], X_tr.shape[3]
        T = X_tr.shape[1]
        H = y_tr.shape[1]

        # Model-specific tensor permutations
        if model_name == "STGCN":
            X_tr = X_tr.permute(0, 3, 1, 2).contiguous()  # (B, C, T, N)
            X_te = X_te.permute(0, 3, 1, 2).contiguous()
        elif model_name == "GraphWaveNet":
            X_tr = X_tr.permute(0, 3, 2, 1).contiguous()  # (B, C, N, T)
            X_te = X_te.permute(0, 3, 2, 1).contiguous()
            y_tr = y_tr.unsqueeze(-1)  # GraphWaveNet outputs (B, 1, N, T)
            y_te = y_te.unsqueeze(-1)
        elif model_name == "ASTGCN":
            X_tr = X_tr.permute(0, 2, 3, 1).contiguous()  # (B, N, C, T)
            X_te = X_te.permute(0, 2, 3, 1).contiguous()
            y_tr = y_tr.permute(0, 2, 1).contiguous()  # (B, N, T)
            y_te = y_te.permute(0, 2, 1).contiguous()

        data_shapes = {"N": N, "C": C, "T": T, "H": H}

    # Build model
    model = build_model(model_name, dataset_name, device, data_shapes, adj)

    # Determine batch_size
    defaults = MODEL_DEFAULTS[model_name]
    batch_size = args.batch_size or defaults["batch_size"]

    if mode in ("train", "both"):
        dataset = TensorDataset(X_tr, y_tr)
        # TLAST rebuttal used 10% val; other baselines use 20%
        val_ratio = 0.1 if model_name == "TLAST" else 0.2
        n_val = max(1, int(val_ratio * len(dataset)))
        n_train = len(dataset) - n_val
        train_set, val_set = random_split(dataset, [n_train, n_val])
        train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True)
        val_loader = DataLoader(val_set, batch_size=batch_size, shuffle=False)

        print(f"  Training samples: {n_train}, Validation samples: {n_val}")
        print(f"  Checkpoints will be saved to: {train_ckpt_path}")
        train_model(model, model_name, train_loader, val_loader, device, train_ckpt_path,
                    data_shapes, epochs=args.epochs, lr=args.lr, patience=args.patience)

    if mode in ("test", "both"):
        load_path = train_ckpt_path if mode == "both" else test_ckpt_path
        if not os.path.exists(load_path):
            print(f"  [ERROR] Checkpoint not found: {load_path}")
            return
        print(f"  Loading checkpoint from: {load_path}")
        # For evaluation, use the original y_test (numpy) for proper inverse transform
        if model_name in ("SATEformer", "GraphWaveNet", "ASTGCN", "iTransformer", "TLAST"):
            y_test_for_eval = y_te.cpu().numpy()
        else:
            y_test_for_eval = y_test

        evaluate_model(model, model_name, X_te, y_test_for_eval, device,
                       load_path, data_shapes)


def main():
    parser = argparse.ArgumentParser(
        description="Unified baseline training & evaluation for traffic speed prediction",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__
    )
    parser.add_argument("--model", type=str, required=True,
                        choices=ALL_MODELS + ["all"],
                        help="Model to run (or 'all' for all models)")
    parser.add_argument("--dataset", type=str, required=True,
                        choices=list(DATASET_CONFIG.keys()) + ["all"],
                        help="Dataset to use (or 'all' for all datasets)")
    parser.add_argument("--mode", type=str, default="both",
                        choices=["train", "test", "both"],
                        help="Run mode: train, test, or both (default: both)")
    parser.add_argument("--save_dir", type=str, default="checkpoints/baselines",
                        help="Directory to save model checkpoints during training (default: checkpoints/baselines)")
    parser.add_argument("--load_dir", type=str, default="checkpoints/baselines",
                        help="Directory to load model checkpoints for testing (default: checkpoints/baselines)")
    parser.add_argument("--data_dir", type=str, default=None,
                        help="Repo root if data/ is not relative to cwd (default: use data/ under cwd)")
    parser.add_argument("--epochs", type=int, default=100,
                        help="Max training epochs (default: 100)")
    parser.add_argument("--batch_size", type=int, default=None,
                        help="Batch size (default: model-specific)")
    parser.add_argument("--lr", type=float, default=None,
                        help="Learning rate (default: model-specific)")
    parser.add_argument("--patience", type=int, default=None,
                        help="Early stopping patience (default: model-specific)")

    args = parser.parse_args()

    models = ALL_MODELS if args.model == "all" else [args.model]
    datasets = list(DATASET_CONFIG.keys()) if args.dataset == "all" else [args.dataset]

    for ds in datasets:
        for m in models:
            try:
                run_single(m, ds, args.mode, args)
            except Exception as e:
                print(f"\n  [ERROR] {m} on {ds}: {e}")
                import traceback
                traceback.print_exc()


if __name__ == "__main__":
    main()