"""
CW3: Entry Point for MLP (PROVIDED — do not modify).

Usage:
    # Train with default settings
    python run.py

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

    # Custom architecture
    python run.py --hidden_dims 256 128 --activation relu --learning_rate 0.01

    # Deeper network
    python run.py --hidden_dims 256 128 64 --momentum 0.9

    # Hyperparameter sweeps (for report)
    python run.py --sweep hidden_dim
    python run.py --sweep num_layers
    python run.py --sweep learning_rate
"""

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, sweep_hyperparameter


# Predefined sweep configurations
SWEEP_CONFIGS = {
    "hidden_dim": {
        "param_name": "hidden_dim",
        "param_values": [32, 64, 128, 256, 512],
    },
    "num_layers": {
        "param_name": "num_layers",
        "param_values": [1, 2, 3, 4],
    },
    "learning_rate": {
        "param_name": "learning_rate",
        "param_values": [0.001, 0.005, 0.01, 0.05, 0.1],
    },
}


def parse_args():
    parser = argparse.ArgumentParser(
        description="CW3: MLP on FashionMNIST")

    # Mode
    parser.add_argument("--quick", action="store_true",
                        help="Quick mode: smaller dataset, fewer epochs")
    parser.add_argument("--sweep", type=str, default=None,
                        choices=list(SWEEP_CONFIGS.keys()),
                        help="Run a hyperparameter sweep")

    # Architecture
    parser.add_argument("--hidden_dims", type=int, nargs="+", default=None,
                        help="Hidden layer dimensions (e.g., --hidden_dims 256 128)")
    parser.add_argument("--activation", type=str, default=None,
                        choices=["relu", "sigmoid"])

    # Training
    parser.add_argument("--learning_rate", type=float, default=None)
    parser.add_argument("--momentum", type=float, default=None)
    parser.add_argument("--weight_decay", type=float, default=None)
    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)

    # 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 key != "sweep" and hasattr(config, key):
            setattr(config, key, val)

    if args.quick:
        config.quick = True

    if args.sweep:
        sweep_cfg = SWEEP_CONFIGS[args.sweep]
        sweep_hyperparameter(config, **sweep_cfg)
    else:
        train(config)


if __name__ == "__main__":
    main()
