from model import *
from utils import *
import gol
import sys
import argparse
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.optim.lr_scheduler import CosineAnnealingLR
import os
import os.path as osp
from torch.utils.tensorboard import SummaryWriter

def quantize_aware_training(model, device, train_loader, optimizer, epoch):
    lossLayer = torch.nn.CrossEntropyLoss()
    #统计loss和每个参数的grad
    #初始化
    loss_sum = 0.
    grad_dict = {}
     # 可以遍历各种weight和bias 
    for name,param in model.named_parameters():
        grad_dict[name] = torch.zeros_like(param) #param.grad和param形状相同
    
    for batch_idx, (data, target) in enumerate(train_loader, 1):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model.quantize_forward(data)
        # 对一批数据求得的loss是平均值
        loss = lossLayer(output, target)
        loss.backward()
        
        #loss和grads累加
        loss_sum += loss
        for name,param in model.named_parameters():
            if param.grad is not None:
                # print('-------'+name+'-------')
                grad_dict[name] += param.grad
                # print(grad_dict[name])

        # print(grad_dict.items())
        # input()
        optimizer.step()

        if batch_idx % 50 == 0:
            print('Quantize Aware Training Epoch: {} [{}/{}]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset), loss.item()
            ))
    
    batch_size = len(train_loader.batch_sampler)
    #对不同batch累加值求平均
    for name,grad in grad_dict.items():
        grad_dict[name] = grad / batch_size
    loss_avg = loss_sum / batch_size
    return loss_avg, grad_dict

def full_inference(model, test_loader):
    correct = 0
    for i, (data, target) in enumerate(test_loader, 1):
        data, target = data.to(device), target.to(device)
        output = model(data)
        pred = output.argmax(dim=1, keepdim=True)
        correct += pred.eq(target.view_as(pred)).sum().item()
    print('\nTest set: Full Model Accuracy: {:.2f}%\n'.format(100. * correct / len(test_loader.dataset)))

def train(model, device, train_loader, optimizer, epoch):
    model.train()
    lossLayer = torch.nn.CrossEntropyLoss()
    #统计loss和每个参数的grad
    #初始化
    loss_sum = 0.
    grad_dict = {}
    for name,param in model.named_parameters():
        grad_dict[name] = torch.zeros_like(param) #param.grad和param形状相同
    
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = lossLayer(output, target)
        loss.backward()
        
        #loss和grads累加
        loss_sum += loss
        for name,param in model.named_parameters():
            if param.grad is not None:
                # print('-------'+name+'-------')
                grad_dict[name] += param.grad   # 对batch的累加
                # print(grad_dict[name])        

        optimizer.step()

        if batch_idx % 50 == 0:
            print('Train Epoch: {} [{}/{}]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset), loss.item()
            ))

    batch_size = len(train_loader.batch_sampler)  # batch数量
    #对不同batch累加值求平均
    for name,grad in grad_dict.items():
        grad_dict[name] = grad / batch_size   # 对batch的平均
    loss_avg = loss_sum / batch_size
    return loss_avg, grad_dict


def quantize_inference(model, test_loader):
    correct = 0
    for i, (data, target) in enumerate(test_loader, 1):
        data, target = data.to(device), target.to(device)
        output = model.quantize_inference(data)
        pred = output.argmax(dim=1, keepdim=True)
        correct += pred.eq(target.view_as(pred)).sum().item()
    print('\nTest set: Quant Model Accuracy: {:.2f}%\n'.format(100. * correct / len(test_loader.dataset)))


if __name__ == "__main__":

    parser = argparse.ArgumentParser(description='QAT Training')
    parser.add_argument('-m', '--model', metavar='MODEL ARCH', default='resnet18')
    parser.add_argument('-e','--epochs', default=20, type=int, metavar='EPOCHS', help='number of total epochs to run')
    parser.add_argument('-b', '--batch_size', default=128, type=int, metavar='BATCH SIZE', help='mini-batch size (default: 128)')
    parser.add_argument('-j','--workers', default=1, type=int, metavar='WORKERS',help='number of data loading workers (default: 4)')
    parser.add_argument('-lr', '--learning-rate', default=0.001, type=float, metavar='LR', help='initial learning rate', dest='lr')
    parser.add_argument('-wd','--weight_decay',default=0.0001,type=float,metavar='WD',help='lr schduler weight decay',dest='wd')
    parser.add_argument('-t', '--test', dest='test', action='store_true', help='test model on test set')
    args = parser.parse_args()

    batch_size = args.batch_size
    seed = 1
    epochs = args.epochs
    lr = args.lr
    # momentum = 0.5
    weight_decay = args.wd

    torch.manual_seed(seed)

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(device)

    writer = SummaryWriter(log_dir='log/' + args.model  +  '/qat')
    train_loader = torch.utils.data.DataLoader(
        datasets.CIFAR10('../../project/p/data', train=True, download=True,
                       transform=transforms.Compose([
                            transforms.ToTensor(),
                            transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
                       ])),
        batch_size=batch_size, shuffle=True, num_workers=args.workers, pin_memory=False
    )

    test_loader = torch.utils.data.DataLoader(
        datasets.CIFAR10('../../project/p/data', train=False, transform=transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
        ])),
        batch_size=batch_size, shuffle=True, num_workers=args.workers, pin_memory=False
    )


    if args.model == 'ResNet18':
        model = resnet18()
    elif args.model == 'ResNet50':
        model = resnet50()
    elif args.model == 'ResNet152':
        model = resnet152()

    
    writer = SummaryWriter(log_dir='log/' + args.model  +  '/qat')


    # full_file = 'ckpt/cifar10_' + args.model + '.pt'


    # model.load_state_dict(torch.load(full_file))
    model.to(device)

    # optimizer = optim.SGD(model.parameters(), lr=lr, momentum=momentum)
    optimizer = optim.Adam(model.parameters(), lr=lr)
    lr_scheduler = CosineAnnealingLR(optimizer, T_max=epochs)

    load_qat = False
    ckpt_prefix = 'ckpt/qat/'+ args.model + '/'

    loss_sum = 0.
    grad_dict_sum = {}
    grad_dict_avg = {}
    for name,param in model.named_parameters():
        grad_dict_sum[name] = torch.zeros_like(param)
        grad_dict_avg[name] = torch.zeros_like(param)
    # full precision from scratch
    for epoch in range(1, epochs+1):
        # 训练原模型，获取梯度分布
        loss,grad_dict = train(model, device, train_loader, optimizer, epoch)
        # print('loss:%f' % loss_avg)
        writer.add_scalar('Full.loss',loss,epoch)
        for name,grad in grad_dict.items():
            writer.add_histogram('Full.'+name+'_grad',grad,global_step=epoch)

        loss_sum += loss
        loss_avg = loss_sum / epoch        
        for name,grad in grad_dict.items():
            grad_dict_sum[name] += grad_dict[name]  # 对epoch的累加
            grad_dict_avg[name] = grad_dict_sum[name] / epoch  # 对epoch求平均

        ckpt = {
            'epoch'     : epoch,
            'loss'      : loss,
            'loss_sum'  : loss_sum,
            'loss_avg'  : loss_avg,
            'grad_dict_avg' : grad_dict_avg
        }
        if epoch % 5 == 0:
            subdir = 'epoch_%d/' % epoch
            if not os.path.isdir(ckpt_prefix+ subdir):
                os.makedirs(ckpt_prefix+ subdir, mode=0o777)
                os.chmod(ckpt_prefix+ subdir, mode=0o777)
            torch.save(ckpt,ckpt_prefix+ subdir +'full.pt')
        lr_scheduler.step()
    
        
        # loss_avg,grad_dict = quantize_aware_training(model_ptq, device, train_loader, optimizer, epoch)
        # print('qat_loss:%f' % loss_avg)
        # for name,grad in grad_dict.items():
        #     writer.add_histogram('qat_'+name+'_grad',grad,global_step=epoch)


    # QAT from scratch
    quant_type_list = ['INT','POT','FLOAT']
    gol._init()
    for quant_type in quant_type_list:
        num_bit_list = numbit_list(quant_type)

        # 对一个量化类别，只需设置一次bias量化表
        # int由于位宽大，使用量化表开销过大，直接_round即可
        if quant_type != 'INT':
            bias_list = build_bias_list(quant_type)
            gol.set_value(bias_list, is_bias=True)

        for num_bits in num_bit_list:
            e_bit_list = ebit_list(quant_type,num_bits)
            for e_bits in e_bit_list:
                
                if quant_type == 'FLOAT':
                    title = '%s_%d_E%d' % (quant_type, num_bits, e_bits)
                else:
                    title = '%s_%d' % (quant_type, num_bits)

                if load_qat is True and osp.exists(ckpt_prefix+'epoch_20/'+title+'.pt'):
                    continue

                print('\nQAT: '+title)
                
                # model_ptq = AlexNet()
                if args.model == 'ResNet18':
                    model_ptq = resnet18()
                elif args.model == 'ResNet50':
                    model_ptq = resnet50()
                elif args.model == 'ResNet152':
                    model_ptq = resnet152()
                 # 设置量化表

                if quant_type != 'INT':
                    plist = build_list(quant_type, num_bits, e_bits)
                    gol.set_value(plist)
                # model_ptq.load_state_dict(torch.load(full_file))
                model_ptq.to(device)
                model_ptq.quantize(quant_type,num_bits,e_bits)
                model_ptq.train()

                loss_sum = 0.
                grad_dict_sum = {}
                grad_dict_avg = {}
                for name,param in model.named_parameters():
                    grad_dict_sum[name] = torch.zeros_like(param)
                    grad_dict_avg[name] = torch.zeros_like(param)
                for epoch in range(1, epochs+1):
                    loss,grad_dict = quantize_aware_training(model_ptq, device, train_loader, optimizer, epoch)
                    # print('loss:%f' % loss_avg)
                    writer.add_scalar(title+'.loss',loss,epoch)
                    for name,grad in grad_dict.items():
                        writer.add_histogram(title+'.'+name+'_grad',grad,global_step=epoch)

                    loss_sum += loss
                    loss_avg = loss_sum / epoch        
                    for name,param in model.named_parameters():
                        grad_dict_sum[name] += grad_dict[name]
                        grad_dict_avg[name] = grad_dict_sum[name] / epoch

                    ckpt = {
                        'epoch'     : epoch,
                        'loss'      : loss,
                        'loss_sum'  : loss_sum,
                        'loss_avg'  : loss_avg,
                        # 'grad_dict' : grad_dict,
                        # 'grad_dict_sum' : grad_dict_sum,
                        'grad_dict_avg' : grad_dict_avg
                    }
                    if epoch % 5 == 0:
                        subdir = 'epoch_%d/' % epoch
                        if not os.path.isdir(ckpt_prefix+ subdir):
                            os.makedirs(ckpt_prefix+ subdir, mode=0o777)
                            os.chmod(ckpt_prefix+ subdir, mode=0o777)
                        torch.save(ckpt,ckpt_prefix+subdir + title+'.pt')
                    lr_scheduler.step()
    writer.close()
    # # model.eval()
    
    # # full_inference(model, test_loader)

    # num_bits = 8
    # e_bits = 0
    # gol._init()
    # print("qat: INT8")

    # model.quantize('INT',num_bits,e_bits)
    # print('Quantization bit: %d' % num_bits)

    # if load_quant_model_file is not None:
    #     model.load_state_dict(torch.load(load_quant_model_file))
    #     print("Successfully load quantized model %s" % load_quant_model_file)
    # else:
    #     model.train()

    #     for epoch in range(1, epochs+1):
    #         quantize_aware_training(model, device, train_loader, optimizer, epoch)
    #     # for epoch in range(epochs1 + 1, epochs2 + 1):
    #     #     quantize_aware_training(model, device, train_loader, optimizer2, epoch)

    #     model.eval()
    #     # torch.save(model.state_dict(), save_file)

    # model.freeze()
    # # for name, param in model.named_parameters():
    # #     print(name)
    # #     print(param.data)
    # #     print('----------')

    # # for param_tensor, param_value in model.state_dict().items():
    # #     print(param_tensor, "\t", param_value)
    # quantize_inference(model, test_loader)

    



    
