# main.py
import torch
from train import train_ctqw_model

if __name__ == "__main__":
    # Run a single demo experiment
    acc = train_ctqw_model(
        dataset_name='MUTAG',
        epochs=100,
        folds=10,
        time_steps=torch.tensor([1.0, 2.0, 3.0, 4.0]),
        hidden_dim=64,
        lr=1e-3,
        batch_size=1,
        fusion='cat',
        heads=4,
        use_attention_bias=True,
        use_sequence_model=True,
        num_layers=2,
        dropout=0.3,
        earlystop_patience=20,
        device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    )

    print(f"\n✅ Demo completed. Average accuracy: {acc:.4f}")
