import argparse
import os
import sys
import time
import torch
import torch.nn.functional as F
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import utils
import tabulate
import models
import numpy as np
from tensorboardX import SummaryWriter
from torch.utils.data.sampler import SubsetRandomSampler

if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='SGD/SWA training')
    parser.add_argument('--dir', type=str, default='seed-${1}00', 
                        help='training directory (default: None)')
    parser.add_argument('--dataset', type=str, default='CIFAR10',
                        help='dataset name: CIFAR10 or IMAGENET12')
    parser.add_argument('--data_path', type=str, default="./data", metavar='PATH',
                        help='path to datasets location (default: "./data")')
    parser.add_argument('--batch_size', type=int, default=128, metavar='N',
                        help='input batch size (default: 128)')
    parser.add_argument('--num_workers', type=int, default=4, metavar='N',
                        help='number of workers (default: 4)')
    parser.add_argument('--model', type=str, default='VGG7LP', metavar='MODEL',
                        help='model name (default: None)')
    parser.add_argument('--epochs', type=int, default=2, metavar='N',
                        help='number of epochs to train (default: 200)')
    parser.add_argument('--seed', type=int, default=200, metavar='N',
                        help='random seed (default: 1)')
    parser.add_argument('--log-name', type=str, default='', metavar='S',
                        help="Name for the log dir")
    parser.add_argument('--log-error', action='store_true', default=False,
                        help='Whether to log quantization error of weight and grad')
    parser.add_argument('--wl-weight', type=int, default=2, metavar='N',
                        help='word length in bits for weight output; -1 if full precision.')
    parser.add_argument('--fl-weight', type=int, default=-1, metavar='N',
                        help='float length in bits for weight output; -1 if full precision.')
    parser.add_argument('--wl-grad', type=int, default=8, metavar='N',
                        help='word length in bits for gradient; -1 if full precision.')
    parser.add_argument('--fl-grad', type=int, default=-1, metavar='N',
                        help='float length in bits for gradient; -1 if full precision.')
    parser.add_argument('--wl-activate', type=int, default=8, metavar='N',
                        help='word length in bits for layer activattions; -1 if full precision.')
    parser.add_argument('--fl-activate', type=int, default=-1, metavar='N',
                        help='float length in bits for layer activations; -1 if full precision.')
    parser.add_argument('--wl-error', type=int, default=8, metavar='N',
                        help='word length in bits for backward error; -1 if full precision.')
    parser.add_argument('--fl-error', type=int, default=-1, metavar='N',
                        help='float length in bits for backward error; -1 if full precision.')
    parser.add_argument('--wl-rand', type=int, default=16, metavar='N',
                        help='word length in bits for rand number; -1 if full precision.')

    args = parser.parse_args()

    torch.backends.cudnn.enabled = True
    torch.backends.cudnn.benchmark = True
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed(args.seed)
    np.random.seed(args.seed)


    # Tensorboard Writer
    if args.log_name != "":
        log_name = args.log_name
    else:
        log_name = "time%d"%int(time.time())
    print("Logging at {}".format(log_name))
    writer = SummaryWriter(log_dir=os.path.join(".", "runs", log_name))


    # Select quantizer
    def quant_summary(wl, fl):
        if wl == -1:
            return "float"
        else:
            return "wage-{}".format(wl)

    w_summary = quant_summary(args.wl_weight, args.fl_weight)       # wl_weight=2,fk_weight=-1
    g_summary = quant_summary(args.wl_grad, args.fl_grad)           # wl_grad=8.fl_grad=-1
    a_summary = quant_summary(args.wl_activate, args.fl_activate)   # wl_acitvate=8,fl_activate=-1
    e_summary = quant_summary(args.wl_error, args.fl_error)         # wl_error=8,fl_error=-1
    print("W:{}, A:{}, G:{}, E:{}".format(w_summary, a_summary, g_summary, e_summary))

    weight_quantizer = lambda x, scale: models.QW(x, args.wl_weight, scale)
    grad_clip = lambda x : models.C(x, args.wl_weight)
    if args.wl_weight==-1: weight_quantizer = None
    if args.wl_grad ==-1: grad_quantier = None


    dir_name = args.dir + "-seed-{}".format(args.seed)
    print('Preparing checkpoint directory {}'.format(dir_name))
    try:
        os.makedirs(dir_name)
    except:
        pass
    with open(os.path.join(dir_name, 'command.sh'), 'w') as f:
        f.write('python '+' '.join(sys.argv))
        f.write('\n')

    assert args.dataset in ["CIFAR10"]
    print('Loading dataset {} from {}'.format(args.dataset, args.data_path))

    if args.dataset=="CIFAR10":
        ds = getattr(datasets, args.dataset)
        path = os.path.join(args.data_path, args.dataset.lower())
        transform_train = transforms.Compose([
            transforms.RandomCrop(32, padding=4),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
        ])
        # 具体是对图像进行各种转换操作，并用函数compose将这些转换操作组合起来
        transform_test = transforms.Compose([
            transforms.ToTensor(),
        ])
        train_set = ds(path, train=True, download=True, transform=transform_train)
        val_set = ds(path, train=True, download=True, transform=transform_test)
        test_set = ds(path, train=False, download=True, transform=transform_test)
        train_sampler = None
        val_sampler = None
        num_classes = 10

    loaders = {
        'train': torch.utils.data.DataLoader(
            train_set,
            batch_size=args.batch_size,
            shuffle=(train_sampler is None),
            sampler=train_sampler,
            num_workers=args.num_workers,
            pin_memory=True
        ),
        'val': torch.utils.data.DataLoader(
            train_set,
            batch_size=args.batch_size,
            sampler=val_sampler,
            num_workers=args.num_workers,
            pin_memory=True
        ),
        'test': torch.utils.data.DataLoader(
            test_set,
            batch_size=args.batch_size,
            shuffle=False,
            num_workers=args.num_workers,
            pin_memory=True
        )
    }

    # Build model
    print('Model: {}'.format(args.model))
    model_cfg = getattr(models, args.model)
    model_cfg.kwargs.update({
        "wl_activate":args.wl_activate, "fl_activate":args.fl_activate,
        "wl_error":args.wl_error, "fl_error":args.fl_error,
        "wl_weight":args.wl_weight,
    })

    if args.log_error:
        model = model_cfg.base(
                *model_cfg.args,
                num_classes=num_classes, writer=writer,
                **model_cfg.kwargs)
    else:
        model = model_cfg.base(
                *model_cfg.args,
                num_classes=num_classes, writer=None,
                **model_cfg.kwargs)
    model.cuda()
    for name, param_acc in model.weight_acc.items():
        model.weight_acc[name] = param_acc.cuda()

    criterion = utils.SSE

    def schedule(epoch):
        if epoch < 200:
            return 8.0
        elif epoch < 250:
            return 1
        else:
            return 1/8.

    start_epoch = 0

    # Prepare logging
    columns = ['ep', 'lr', 'tr_loss', 'tr_acc', 'tr_acc2',
            'te_loss', 'te_acc', 'te_acc2', 'time']

    def log_result(writer, name, res, step):
        writer.add_scalar("{}/loss".format(name),     res['loss'],            step)
        writer.add_scalar("{}/acc_perc".format(name), res['accuracy'],        step)
        writer.add_scalar("{}/err_perc".format(name), 100. - res['accuracy'], step)

    for epoch in range(start_epoch, args.epochs):
        time_ep = time.time()
        lr = schedule(epoch)
        writer.add_scalar("lr", lr, epoch)
        grad_quantizer = lambda x : models.QG(x, args.wl_grad, args.wl_rand, lr)       
        # wl_grad = 8, wl_rand=16

        train_res = utils.train_epoch(
                loaders['train'], model, criterion,
                weight_quantizer, grad_quantizer, writer, epoch,
                log_error=args.log_error,
                wage_quantize=True,
                wage_grad_clip=grad_clip
        )
        log_result(writer, "train", train_res, epoch+1)

        # Validation
        test_res = utils.eval(loaders['test'], model, criterion, weight_quantizer)
        log_result(writer, "test", test_res, epoch+1)

        time_ep = time.time() - time_ep
        values = [epoch + 1, lr, train_res['loss'], train_res['accuracy'],
                train_res['semi_accuracy'], test_res['loss'], test_res['accuracy'],
                test_res['semi_accuracy'], time_ep]

        table = tabulate.tabulate([values], columns, tablefmt='simple', floatfmt='8.4f')
        if epoch % 20 == 0:
            table = table.split('\n')
            table = '\n'.join([table[1]] + table)
        else:
            table = table.split('\n')[2]
        print(table)


