import sys
import time

import torch
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Variable

from archived.s3 import download_file

from archived.pytorch_model import MobileNet


# lambda setting
merged_bucket = "cnn-updates"
tmp_bucket = "cnn-grads"
weights_prefix = 'w_'
gradients_prefix = 'g_'
local_dir = "/tmp"

# dataset setting
training_file = "training.pt"
test_file = "test.pt"

# sync up mode
sync_mode = 'grad_avg'
#sync_mode = 'model_avg'
sync_step = 1

# learning algorithm setting
learning_rate = 0.01
batch_size = 32
num_epochs = 1


def handler(event, context):
    start_time = time.time()
    bucket = event['data_bucket']
    worker_index = event['rank']
    num_worker = event['num_workers'] 
    key = 'training_{}.pt'.format(worker_index)
    print('data_bucket = {}\n worker_index:{}\n num_worker:{}\n key:{}'.format(bucket, worker_index, num_worker, key))

    # read file from s3
    readS3_start = time.time()
    train_path = download_file(bucket, key)
    test_path = download_file(bucket, test_file)
    train_set = torch.load(train_path)
    test_set= torch.load(test_path)
    print("read data cost {} s".format(time.time() - readS3_start))
    print(train_set)
    train_loader = torch.utils.data.DataLoader(train_set, batch_size=batch_size, shuffle=True)
    print(enumerate(train_loader))
    test_loader = torch.utils.data.DataLoader(test_set, batch_size=100, shuffle=False)
    classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
    device = 'cpu'
    # best_acc = 0  # best test accuracy
    # start_epoch = 0  # start from epoch 0 or last checkpoint epoch
    #Model
    print('==> Building model..')
    # net = VGG('VGG19')
    # net = ResNet18()
    #net = ResNet50()
    # net = PreActResNet18()
    # net = GoogLeNet()
    # net = DenseNet121()
    # net = ResNeXt29_2x64d()
    net = MobileNet()
    # net = MobileNetV2()
    # net = DPN92()
    # net = ShuffleNetG2()
    # net = SENet18()
    # net = ShuffleNetV2(1)
    # net = EfficientNetB0()
    
    print("Model: MobileNet")

    net = net.to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(net.parameters(), lr=learning_rate, momentum=0.9, weight_decay=5e-4)
    
    for epoch in range(num_epochs):
        time_record = train(epoch, net, train_loader, optimizer, criterion, device, worker_index, num_worker, sync_mode, sync_step)
        test(epoch, net, test_loader, criterion, device)
    put_object("time-record-s3", "time_{}".format(worker_index), pickle.dumps(time_record))


# Training
def train(epoch, net, trainloader, optimizer, criterion, device, worker_index, num_worker, sync_mode, sync_step):
    # print('\nEpoch: %d' % epoch)
    net.train()
    train_loss = 0
    correct = 0
    total = 0
    sync_epoch_time = []
    write_local_epoch_time = []
    calculation_epoch_time = []
    
    for batch_idx, (inputs, targets) in enumerate(trainloader):
        print("------worker {} epoch {} batch {} sync mode '{}'------".format(worker_index, epoch+1, batch_idx+1,sync_mode))
        batch_start = time.time()
        
        inputs, targets = inputs.to(device), targets.to(device)
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        tmp_calculation_time = time.time()-batch_start
        print("forward and backward cost {} s".format(tmp_calculation_time))
        if batch_idx != 0:
            calculation_epoch_time.append(tmp_calculation_time)
        if sync_mode == 'grad_avg':
            sync_start = time.time()
            gradients = [param.grad.data.numpy() for param in net.parameters()]
            # print("[Worker {}] Gradients before sync = {}".format(worker_index, gradients[0][0]))

            put_object_start = time.time()
            put_object(tmp_bucket, gradients_prefix + str(worker_index), pickle.dumps(gradients))
            tmp_write_local_epoch_time = time.time() - put_object_start
            print("write local gradients cost {} s".format(tmp_write_local_epoch_time))
            if batch_idx !=0 :
                write_local_epoch_time.append(tmp_write_local_epoch_time)
            file_postfix = "{}_{}".format(epoch, batch_idx)
            if worker_index == 0:
                # merge all workers
                merged_value_start = time.time()
                merged_value = merge_all_workers(tmp_bucket, num_worker, gradients_prefix)
                print("merged_value cost {} s".format(time.time() - merged_value_start))
                

                print("size of gradients = {}".format(sys.getsizeof(pickle.dumps(merged_value))/1024/1024))
                
                put_merged_start = time.time()
                # upload merged value to S3
                put_merged(merged_bucket, merged_value, gradients_prefix, file_postfix)                  
                print("put_merged cost {} s".format(time.time() - put_merged_start))                   
                delete_expired(merged_bucket, epoch, batch_idx, gradients_prefix)
            else:
                read_merged_start = time.time()
                # get merged value from redis
                merged_value = get_merged(merged_bucket, gradients_prefix, file_postfix)
                print("read_merged cost {} s".format(time.time() - read_merged_start))

            for layer_index, param in enumerate(net.parameters()):
                param.grad = Variable(torch.from_numpy(merged_value[layer_index]))

            tmp_sync_time = time.time() - sync_start
            print("synchronization cost {} s".format(tmp_sync_time))
            if batch_idx != 0:
                sync_epoch_time.append(tmp_sync_time)

            optimizer.step()

        train_loss += loss.item()
        _, predicted = outputs.max(1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()
        print("batch cost {} s".format(time.time() - batch_start))
        if (batch_idx + 1) % 10 == 0:
            print('Epoch: {}, Step: {}, Loss:{}'.format(epoch+1, batch_idx+1, loss.data))
    return sync_epoch_time, write_local_epoch_time, calculation_epoch_time

def test(epoch, net, testloader, criterion, device):
    # global best_acc
    net.eval()
    test_loss = 0
    correct = 0
    total = 0
    with torch.no_grad():
        for batch_idx, (inputs, targets) in enumerate(testloader):
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = net(inputs)
            loss = criterion(outputs, targets)

            test_loss += loss.item()
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()

            # progress_bar(batch_idx, len(testloader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)'
            #     % (test_loss/(batch_idx+1), 100.*correct/total, correct, total))

    # Save checkpoint.
    acc = 100.*correct/total
    print("Accuracy of epoch {} on test set is {}".format(epoch, acc))