import os
import torch.nn as nn
import numpy as np
import argparse
import dq.nets as nets
import dq.datasets as datasets
import dq.methods as methods
from torchvision import transforms
from util.utils import *
from datetime import datetime
from time import sleep


def main():
    parser = argparse.ArgumentParser(description='Parameter Processing')

    # Basic arguments
    parser.add_argument('--dataset', type=str, default='CIFAR10', help='dataset')
    parser.add_argument('--model', type=str, default='ResNet18', help='model')
    parser.add_argument('--selection', type=str, default="uniform", help="selection method")
    parser.add_argument('--num_exp', type=int, default=5, help='the number of experiments')
    parser.add_argument('--batch', type=int, default=128, help='the number of batch size for selection')
    parser.add_argument('--gpu', default=None, nargs="+", type=int, help='gpu id to use')
    parser.add_argument('--data_path', type=str, default='data', help='dataset path')
    parser.add_argument('--print_freq', '-p', default=50, type=int, help='print frequency (default: 20)')
    parser.add_argument('--fraction', default=0.1, type=float, help='fraction of data to be selected (default: 0.1)')
    parser.add_argument('--seed', default=int(time.time() * 1000) % 100000, type=int, help="random seed")
    parser.add_argument('-j', '--workers', default=4, type=int, metavar='N',
                        help='number of data loading workers (default: 4)')

    # Selecting
    parser.add_argument("--selection_epochs", "-se", default=40, type=int,
                        help="number of epochs whiling performing selection on full dataset")
    parser.add_argument('--selection_momentum', '-sm', default=0.9, type=float, metavar='M',
                        help='momentum whiling performing selection (default: 0.9)')
    parser.add_argument('--selection_weight_decay', '-swd', default=5e-4, type=float,
                        metavar='W', help='weight decay whiling performing selection (default: 5e-4)',
                        dest='selection_weight_decay')
    parser.add_argument('--selection_optimizer', "-so", default="SGD",
                        help='optimizer to use whiling performing selection, e.g. SGD, Adam')
    parser.add_argument("--selection_nesterov", "-sn", default=True, type=str_to_bool,
                        help="if set nesterov whiling performing selection")
    parser.add_argument('--selection_lr', '-slr', type=float, default=0.1, help='learning rate for selection')
    parser.add_argument("--selection_test_interval", '-sti', default=1, type=int, help=
    "the number of training epochs to be preformed between two test epochs during selection (default: 1)")
    parser.add_argument("--selection_test_fraction", '-stf', type=float, default=1.,
             help="proportion of test dataset used for evaluating the model while preforming selection (default: 1.)")
    parser.add_argument('--balance', default=True, type=str_to_bool,
                        help="whether balance selection is performed per class")

    # Algorithm
    parser.add_argument('--submodular', default="GraphCut", help="specifiy submodular function to use")
    parser.add_argument('--submodular_greedy', default="NaiveGreedy", help="specifiy greedy algorithm for submodular optimization")

    # Checkpoint and resumption
    parser.add_argument('--save_path', "-sp", type=str, default='', help='path to save results (default: do not save)')
    parser.add_argument('--replace', action='store_true', default=False)
    parser.add_argument('--pretrained', action='store_true', default=False)

    args = parser.parse_args()
    args.device = 'cuda' if torch.cuda.is_available() else 'cpu'

    args.selection_batch = args.batch
    if args.save_path != "" and not os.path.exists(args.save_path):
        os.mkdir(args.save_path)
    if not os.path.exists(args.data_path):
        os.mkdir(args.data_path)
    checkpoint = {}
    start_exp = 0
    start_epoch = 0

    for exp in range(start_exp, args.num_exp):
        print('\n================== Exp %d ==================\n' % exp)
        print(args)
        channel, im_size, num_classes, class_names, mean, std, dst_train, dst_test = datasets.__dict__[args.dataset] \
            (args.data_path)
        args.channel, args.im_size, args.num_classes, args.class_names = channel, im_size, num_classes, class_names

        torch.random.manual_seed(args.seed)

        if exp == 0:
            avail_indices = np.arange(len(dst_train))
            coreset_size = round(len(dst_train) * args.fraction)

        selection_args = dict(epochs=args.selection_epochs,
                              balance=args.balance,
                              greedy=args.submodular_greedy,
                              function=args.submodular,
                              torchvision_pretrain=args.pretrained
                              )
        dst_train = torch.utils.data.Subset(dst_train, avail_indices)
        fraction = coreset_size / len(avail_indices)
        print('Exp: {}, Dst Size: {}, Fraction: {}'.format(exp, len(dst_train), fraction))
        method = methods.__dict__[args.selection](dst_train, args, fraction, args.seed, **selection_args)
        subset = method.select()
        mapped_indices = avail_indices[subset["indices"]]
        print('Exp: {}, Available indices: {}, Select Subset: {}'.format(exp, len(avail_indices), len(mapped_indices)))

        select_save_path = os.path.join(
            args.save_path, 'select_indices_{}_exp{}.npy'.format(args.dataset, exp)
        )
        np.save(select_save_path, mapped_indices)
        if not args.replace:
            avail_indices = np.delete(avail_indices, subset["indices"])

        if not args.replace and exp == args.num_exp - 2 and abs(args.fraction * args.num_exp - 1) < 1e-5:
            select_save_path = os.path.join(
                args.save_path, 'select_indices_{}_exp{}.npy'.format(args.dataset, exp + 1)
            )
            np.save(select_save_path, avail_indices)
            return
        else:
            continue


if __name__ == '__main__':
    main()
