"""
CW4: Entry Point for CNN (PROVIDED — do not modify).

Usage:
    # Train SimpleCNN with default settings
    python run.py

    # Quick mode for debugging
    python run.py --quick

    # Train with data augmentation
    python run.py --use_augmentation

    # Train DeepCNN (bonus)
    python run.py --model_type deep_cnn --num_epochs 20

    # Custom hyperparameters
    python run.py --learning_rate 0.0005 --optimizer_type sgd --batch_size 32
"""

import argparse
import sys
import os

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

from config import Config
from trainer import train


def parse_args():
    parser = argparse.ArgumentParser(
        description="CW4: CNN on FashionMNIST")

    # Mode
    parser.add_argument("--quick", action="store_true",
                        help="Quick mode: fewer epochs")

    # Model
    parser.add_argument("--model_type", type=str, default=None,
                        choices=["simple_cnn", "deep_cnn"])

    # Training
    parser.add_argument("--learning_rate", type=float, default=None)
    parser.add_argument("--optimizer_type", type=str, default=None,
                        choices=["adam", "sgd"])
    parser.add_argument("--batch_size", type=int, default=None)
    parser.add_argument("--num_epochs", type=int, default=None)
    parser.add_argument("--seed", type=int, default=None)

    # Augmentation
    parser.add_argument("--use_augmentation", action="store_true")

    # Data
    parser.add_argument("--data_dir", type=str, default=None)
    parser.add_argument("--output_dir", type=str, default=None)

    return parser.parse_args()


def main():
    args = parse_args()
    config = Config()

    # Override config with command-line arguments
    for key, val in vars(args).items():
        if val is not None and hasattr(config, key):
            setattr(config, key, val)

    if args.quick:
        config.quick = True
    if args.use_augmentation:
        config.use_augmentation = True

    train(config)


if __name__ == "__main__":
    main()
