import arg_parser
import utils
import datasets
import models
import pruner
import attacks
import os

import training_utils


def main():
    args = arg_parser.parse_args_victim_training()
    utils.set_seed(args.seed)

    if args.dataset == 'cifar10':
        loader = datasets.CIFAR10(
            dir=args.dataset_dir, ffcv_dir=args.ffcv_dir, batch_size=args.batch_size)
    elif args.dataset == 'cifar100':
        loader = datasets.CIFAR100(
            dir=args.dataset_dir, ffcv_dir=args.ffcv_dir, batch_size=args.batch_size)
    elif args.dataset == 'tinyimagenet':
        loader = datasets.TinyImageNet(
            dir=args.dataset_dir, ffcv_dir=args.ffcv_dir, batch_size=args.batch_size)
    elif args.dataset == 'mnist':
        loader = datasets.MNIST(
            dir=args.dataset_dir, batch_size=args.batch_size)
    else:
        raise NotImplementedError(f"Dataset {args.dataset} not implemented!")
    model = models.get_model(args.arch, args)
    model = model.cuda()

    if args.dataset != 'mnist':
        prefix = training_utils.get_model_name(
            seed=args.seed, 
            kernel_size=args.kernel_size, 
            activation_function=args.act_func,
            pruning_ratio=args.pruning_ratio,
            struct=args.structured_pruning,
            robust=args.robust_train) + '_'
    else:
        prefix = "seed{}_conv{}_fc{}_kernel{}_act{}_prune{}_".format(
            args.seed, args.num_conv, args.num_fc, args.kernel_size, args.act_func, args.pruning_ratio)
        if args.structured_pruning:
            prefix += "struct_"
        if args.robust_train:
            prefix += "robust_"

    pruner.omp(model, loader, args, prefix=prefix)

    if args.attack is not None:
        atk = attacks.get_attack_normalized(model, args.attack, args)
        prefix = prefix[:-1]
        attack_dir = os.path.join(args.attack_save_dir, prefix)
        attacks.generate_attack_images(
            model, loader, atk, attack_dir)


if __name__ == "__main__":
    main()
