"""
Main entry point for CTM-DQN speed limit control system.
"""
import argparse
from train import train
from test import test


def main():
    parser = argparse.ArgumentParser(
        description="DQN-based Dynamic Speed Limit Control with CTM"
    )
    parser.add_argument(
        "--mode",
        type=str,
        choices=["train", "test"],
        default="train",
        help="Mode: train or test",
    )
    parser.add_argument(
        "--config",
        type=str,
        default="config.yaml",
        help="Path to configuration file",
    )
    parser.add_argument(
        "--model",
        type=str,
        default=None,
        help="Path to model checkpoint (for testing)",
    )

    args = parser.parse_args()

    if args.mode == "train":
        print("Starting training mode...")
        train(args.config)
    elif args.mode == "test":
        print("Starting testing mode...")
        test(args.config, args.model)


if __name__ == "__main__":
    main()