"""
ml/training/train_sequence_vae.py
Train the 1D-CNN Sequence VAE on normal track geometry sequences (tc.v1 SOTA).
"""

from ml.scripts.train_sequence_vae import train_sequence_vae

__all__ = ["train_sequence_vae"]

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Train 1D-CNN Sequence VAE for Novel Geometry.")
    parser.add_argument("--epochs", type=int, default=25, help="Number of training epochs")
    parser.add_argument("--beta", type=float, default=0.01, help="Beta-VAE KL weight")
    parser.add_argument("--batch-size", type=int, default=64, help="Batch size")
    parser.add_argument("--lr", type=float, default=0.001, help="Learning rate")
    parser.add_argument("--save-path", default="artifacts/checkpoints/geometry/sequence_vae.pt")
    args = parser.parse_args()

    train_sequence_vae(
        epochs=args.epochs,
        beta=args.beta,
        batch_size=args.batch_size,
        lr=args.lr,
        save_path=args.save_path,
    )