from core.datasets.factory import create_dataset
from core.models.BaseModel import BaseModel
from scripts.utils import set_random_seed, DATASETS

import argparse


def main():
    args = parse_arguments()

    set_random_seed(args.seed)

    dataset = create_dataset(args.dataset, args.custom_dim, args.seed)
    (x_train, y_train), (x_test, y_test) = dataset.get_data()

    model = BaseModel(dataset, args.seed)

    model.build(
        hidden_layers=args.hidden_layers,
        activation=args.activation,
        learning_rate=args.learning_rate
    )

    model.train(
        x=x_train,
        y=y_train,
        batch_size=args.batch_size,
        epochs=args.epochs,
        patience=args.patience
    )

    score = model.evaluate(
        x=x_test,
        y=y_test,
        batch_size=args.batch_size
    )

    print("Test loss:", score[0])
    print("Test accuracy:", score[1])

    model.save(model.path)
    print(f"Model saved as {model.path}")


def parse_arguments():
    parser = argparse.ArgumentParser(description='Train a neural network model',
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    # Dataset selection
    parser.add_argument('--dataset', type=str, required=True, choices=DATASETS,
                        help='Dataset to use for training')

    # Custom dataset parameters
    parser.add_argument('--custom-dim', type=int, default=0,
                        help='Dimensionality of the custom dataset')

    # Model architecture parameters
    parser.add_argument('--hidden-layers', nargs='+', type=int, default=[10, 10],
                        help='Hidden layer sizes')
    parser.add_argument('--activation', type=str, default='relu',
                        help='Activation function')

    # Training parameters
    parser.add_argument('--learning-rate', type=float, default=0.001,
                        help='Learning rate')
    parser.add_argument('--batch-size', type=int, default=16,
                        help='Batch size')
    parser.add_argument('--epochs', type=int, default=500,
                        help='Number of epochs')
    parser.add_argument('--patience', type=int, default=10,
                        help='Early stopping patience')
    parser.add_argument('--seed', type=int, default=0,
                        help='Random seed for reproducibility')

    args = parser.parse_args()

    if args.dataset == 'custom' and args.custom_dim == 0:
        parser.error("Custom dataset requires --custom_dim to be set.")

    if args.dataset != 'custom' and args.custom_dim != 0:
        parser.error("Custom dataset parameters are only applicable when --dataset is set to 'custom'.")

    return args


if __name__ == "__main__":
    main()
