from core.datasets.factory import create_dataset
from core.models.factory import create_model
from scripts.utils import DATASETS

import argparse


def main():
    args = parse_arguments()

    dataset = create_dataset(args.dataset, args.custom_dim, args.seed)

    sample = dataset.get_samples(
        label=args.label,
        num_samples=1
    )[0]

    model = create_model(dataset, args.seed)
    predicted_class = model.classify([sample])

    print(f'Predicted class: {predicted_class}')


def parse_arguments():
    parser = argparse.ArgumentParser(description='Classify an image or sample from a dataset',
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    # Dataset parameters
    parser.add_argument('--dataset', type=str, required=True, choices=DATASETS,
                        help='Dataset to use')
    parser.add_argument('--seed', type=int, default=0,
                        help='Seed for the dataset')

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

    # Classification parameters
    parser.add_argument('--label', type=int, default=0,
                        help='Label to classify')

    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()
