"""
CW2: Entry Point for Logistic Regression (PROVIDED — do not modify).

Usage:
    # Train with default settings (cross-entropy loss)
    python run.py

    # Train with a specific loss function
    python run.py --loss_type hinge
    python run.py --loss_type exponential
    python run.py --loss_type squared

    # Quick mode for debugging (smaller dataset, fewer epochs)
    python run.py --quick

    # Compare all four loss functions (generates comparison plot)
    python run.py --compare_all

    # Custom hyperparameters
    python run.py --learning_rate 0.01 --batch_size 64 --num_epochs 50
"""

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


def parse_args():
    parser = argparse.ArgumentParser(
        description="CW2: Logistic Regression on FashionMNIST")

    # Mode
    parser.add_argument("--compare_all", action="store_true",
                        help="Train with all 4 loss functions and compare")
    parser.add_argument("--quick", action="store_true",
                        help="Quick mode: use 10K samples and fewer epochs")

    # Loss
    parser.add_argument("--loss_type", type=str, default=None,
                        choices=["cross_entropy", "hinge", "exponential", "squared"],
                        help="Loss function to use (default: cross_entropy)")

    # Training
    parser.add_argument("--learning_rate", 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)

    # Regularization
    parser.add_argument("--regularization", type=float, default=None,
                        help="L2 regularization strength (bonus)")

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

    if args.quick:
        config.quick = True

    if args.compare_all:
        compare_all_losses(config)
    else:
        train(config)


if __name__ == "__main__":
    main()
