"""
Train Async-VFL on ModelNet-10 dataset
"""

import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
from torch.utils.data import DataLoader
from torch.autograd import Variable

import torchvision.transforms as transforms

import argparse
import numpy as np
import time
import os
import copy
import random
import pickle
import math
import subprocess as sp
from tqdm import tqdm

from scipy.optimize import minimize
from scipy.optimize import Bounds
from scipy.optimize import NonlinearConstraint
from scipy.optimize import BFGS

from models.resnet import *
from models.mvcnn import *
from models.mvcnn_top import *
from models.mvcnn_bottom import *
import util
#from logger import Logger
from custom_dataset import MultiViewDataSet
import sys

MVCNN = 'mvcnn'
RESNET = 'resnet'
MODELS = [RESNET,MVCNN]

def get_gpu_memory():
    """
    Utility function for checking GPU memory
    """
    _output_to_list = lambda x: x.decode('ascii').split('\n')[:-1]

    ACCEPTABLE_AVAILABLE_MEMORY = 1024
    COMMAND = "nvidia-smi --query-gpu=memory.free --format=csv"
    memory_free_info = _output_to_list(sp.check_output(COMMAND.split()))[1:]
    memory_free_values = [int(x.split()[0]) for i, x in enumerate(memory_free_info)]
    print(memory_free_values)
    return memory_free_values


# Set up input arguments
num_clients = int(sys.argv[3])

parser = argparse.ArgumentParser(description='MVCNN-PyTorch')
parser.add_argument('data', metavar='DIR', help='path to dataset')
parser.add_argument('--num_clients', type=int, help='Number of clients to split data between vertically',
                        default=2)
parser.add_argument('--depth', choices=[18, 34, 50, 101, 152], type=int, metavar='N', default=18, help='resnet depth (default: resnet18)')
parser.add_argument('--model', '-m', metavar='MODEL', default=RESNET, choices=MODELS,
                    help='pretrained model: ' + ' | '.join(MODELS) + ' (default: {})'.format(RESNET))
parser.add_argument('--epochs', default=100, type=int, metavar='N', help='number of total epochs to run (default: 100)')
parser.add_argument('-b', '--batch-size', default=4, type=int,
                    metavar='N', help='mini-batch size (default: 4)')
parser.add_argument('--lr', '--learning-rate', default=0.0001, type=float,
                    metavar='LR', help='initial learning rate (default: 0.0001)')
parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
                    help='momentum (default: 0.9)')
parser.add_argument('--lr-decay-freq', default=30, type=float,
                    metavar='W', help='learning rate decay (default: 30)')
parser.add_argument('--lr-decay', default=0.1, type=float,
                    metavar='W', help='learning rate decay (default: 0.1)')
parser.add_argument('--print-freq', '-p', default=10, type=int,
                    metavar='N', help='print frequency (default: 10)')
parser.add_argument('-r', '--resume', default='', type=str, metavar='PATH',
                    help='path to latest checkpoint (default: none)')
parser.add_argument('--pretrained', dest='pretrained', action='store_true', help='use pre-trained model')
parser.add_argument('--prob', type=float, nargs=num_clients+1, default=1,
                        help='Indicates probability distribution to use for workers.')
parser.add_argument('--lr_type', type=str, nargs='?', default="hetero",
                        help='Indicates if learning rates should be heterogeneous or homogeneous.')
parser.add_argument('--priori', type=str, nargs='?', default="a",
                        help='Adjust learning rate adaptively or not')
parser.add_argument('--server_time', type=int, nargs='?', default=1,
                        help='Number of local iteration time it takes for server to communicate.')

# Parse input arguments
args = parser.parse_args()

print('Loading data')

transform = transforms.Compose([
    transforms.CenterCrop(500),
    transforms.Resize(224),
    transforms.ToTensor(),
])


# Load dataset
dset_train = MultiViewDataSet(args.data, 'train', transform=transform)
indices = torch.randperm(len(dset_train))
dset_train_sub = torch.utils.data.Subset(dset_train, indices[:int(len(dset_train)/4)])
train_loader = DataLoader(dset_train_sub, batch_size=args.batch_size, shuffle=True, num_workers=1)

dset_val = MultiViewDataSet(args.data, 'test', transform=transform)
test_loader = DataLoader(dset_val, batch_size=args.batch_size, shuffle=True, num_workers=1)

classes = dset_train.classes
print(len(classes), classes)

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
models = []
optimizers = []
# Make models for each client
for i in range(num_clients+1):
    if i == num_clients:
        model = mvcnn_top(pretrained=args.pretrained,num_classes=len(classes))
    else:
        model = mvcnn_bottom(pretrained=args.pretrained,num_classes=len(classes))

    cudnn.benchmark = True

    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
    models.append(model)
    optimizers.append(optimizer)

# Client views of the server model
server_models = [] 
for i in range(num_clients):
    server_models.append(copy.deepcopy(models[-1]))

# Loss and Optimizer
n_epochs = args.epochs
criterion = nn.CrossEntropyLoss()
coords_per = int(12/num_clients)

best_acc = 0.0
best_loss = 0.0
start_epoch = 0

# Validation and Testing
def eval(models, data_loader):
    """
    Calculate loss and accuracy for a given data_loader
    """
    total = 0.0
    correct = 0.0

    total_loss = 0.0
    n = 0

    for i, (inputs, targets) in enumerate(data_loader):
        with torch.no_grad():
            # Convert from list of 3D to 4D
            inputs = np.stack(inputs, axis=1)

            inputs = torch.from_numpy(inputs)

            #inputs, targets = inputs.cuda(device), targets.cuda(device)
            inputs, targets = Variable(inputs), Variable(targets)

            # Get current embeddings
            H_new = [None] * num_clients
            for i in range(num_clients):
                x_local = inputs[:,coords_per*i:coords_per*(i+1),:,:,:]
                H_new[i] = models[i](x_local)
            # compute output
            outputs = models[-1](torch.cat(H_new,axis=1))
            loss = criterion(outputs, targets)

            total_loss += loss
            n += 1

            _, predicted = torch.max(outputs.data, 1)
            total += targets.size(0)
            correct += (predicted.cpu() == targets.cpu()).sum()

    avg_test_acc = 100 * correct / total
    avg_loss = total_loss / n

    return avg_test_acc, avg_loss


losses = []
accs_train = []
accs_test = []
orig_prob = copy.deepcopy(args.prob)
prob = args.prob

# Get embeddings for all batches from each client
server_H = np.empty((len(train_loader), num_clients), dtype=object)
client_H = []
for i in range(num_clients):
    client_H.append(np.empty((len(train_loader), num_clients), dtype=object))
for step, (inputs, targets) in tqdm(enumerate(train_loader)):
    # Convert from list of 3D to 4D
    inputs = np.stack(inputs, axis=1)
    inputs = torch.from_numpy(inputs)
    #inputs, targets = inputs.cuda(device), targets.cuda(device)
    inputs, targets = Variable(inputs), Variable(targets)
    # Train for each client 
    for i in range(num_clients):
        x_local = inputs[:,coords_per*i:coords_per*(i+1),:,:,:]
        server_H[step,i] = models[i](x_local)
    for k in range(num_clients):
        client_H[k][step] = server_H[step].copy()

train_size = len(train_loader)
server_model = models[-1]
server_optimizer = optimizers[-1]
#server_model.to(device)
#for k in range(num_clients):
#    models[k].to(device)
client_busy = [0] * num_clients
server_busy = 0
queue = {}

# Get initial loss/accuracy
avg_train_acc, avg_loss = eval(models, train_loader)
#avg_test_acc, _ = eval(models, test_loader)

losses.append(avg_loss)
accs_train.append(avg_train_acc)
#accs_test.append(avg_test_acc)

pickle.dump(losses, open(f'loss_mvcnnAsync_NC{args.num_clients}_lr{args.lr_type}_prob{orig_prob}_priori{args.priori}_server{args.server_time}.pkl', 'wb'))
pickle.dump(accs_train, open(f'accs_train_mvcnnAsync_NC{args.num_clients}_lr{args.lr_type}_prob{orig_prob}_priori{args.priori}_server{args.server_time}.pkl', 'wb'))
#pickle.dump(accs_test, open(f'accs_test_mvcnnAsync_NC{args.num_clients}_lr{args.lr_type}_prob{orig_prob}_priori{args.priori}_server{args.server_time}.pkl', 'wb'))


#print('Train Acc: %.2f - Test Acc: %.2f - Loss: %.4f' % (avg_train_acc.item(), avg_test_acc.item(), avg_loss.item()))
print('Train Acc: %.2f - Loss: %.4f' 
        % (avg_train_acc.item(), avg_loss.item()))
# Training / Eval loop
for epoch in range(start_epoch, n_epochs):
    # Set offset for saving loss/accuracy every 6 batches 
    print('\n-----------------------------------')
    print('Epoch: [%d/%d]' % (epoch+1, n_epochs))
    start = time.time()

    for step, (inputs, targets) in enumerate(train_loader):
        # Convert from list of 3D to 4D
        inputs = np.stack(inputs, axis=1)

        inputs = torch.from_numpy(inputs)

        #inputs, targets = inputs.cuda(device), targets.cuda(device)
        inputs, targets = Variable(inputs), Variable(targets)

        # Train clients
        for i in range(num_clients):
            print(client_busy)
            # Is client waiting on the server?
            if i in queue.keys():
                continue
            client_busy[i] -= 1
            if client_busy[i] > 0:
                continue

            # Train for each client with probability
            if random.random() < 1-prob[i]:
                continue

            x_local = inputs[:,coords_per*i:coords_per*(i+1),:,:,:]
            H = client_H[i][step].copy()
            model = models[i]
            optimizer = optimizers[i]

            # Train
            #model.to(device)
            #server_model.to(device)
            outputs = model(x_local)
            H[i] = outputs.cpu()
            Hlist = torch.cat(H.tolist(),axis=1)
            outputs = server_model(Hlist)
            loss = criterion(outputs, targets)

            # compute gradient and do SGD step
            optimizer.zero_grad()
            server_optimizer.zero_grad()
            loss.backward(retain_graph=True)
            optimizer.step()
            #model.to(torch.device("cpu"))
            #server_model.to(torch.device("cpu"))

            # Check if client is ready to comm with server
            H_new = model(x_local).cpu()
            queue[i] = [step, H_new, copy.deepcopy(targets)] 

        # If server is not busy and there are clients in the queue
        if len(queue) > 0:
            # Train server on latest embeddings
            for k in queue:
                k_step = queue[k][0]
                server_H[k_step,k] = queue[k][1]
                y_local = queue[k][2]
                client_busy[k] = args.server_time+len(queue)

                # Train server
                #server_model.to(device)
                # compute output
                Hlist = torch.cat(server_H[k_step].tolist(),axis=1)
                outputs = server_model(Hlist)
                loss = criterion(outputs, y_local)

                # compute gradient and do SGD step
                server_optimizer.zero_grad()
                loss.backward(retain_graph=True)
                server_optimizer.step()
                #server_model.to(torch.device("cpu"))

            # Update client embeddings and server model
            for k in queue:
                client_H[k] = server_H.copy()
                server_models[k] = copy.deepcopy(models[-1])
            queue.clear()

        # Save current loss/accuracy every 15 batches
        if (epoch*63 + step + 1) % 24 == 0:
            avg_train_acc, avg_loss = eval(models, train_loader)
            #avg_test_acc, _ = eval(models, test_loader)

            losses.append(avg_loss)
            accs_train.append(avg_train_acc)
            #accs_test.append(avg_test_acc)
            
            pickle.dump(losses, open(f'loss_mvcnnAsync_NC{args.num_clients}_lr{args.lr_type}_prob{orig_prob}_priori{args.priori}_server{args.server_time}.pkl', 'wb'))
            pickle.dump(accs_train, open(f'accs_train_mvcnnAsync_NC{args.num_clients}_lr{args.lr_type}_prob{orig_prob}_priori{args.priori}_server{args.server_time}.pkl', 'wb'))
            #pickle.dump(accs_test, open(f'accs_test_mvcnnAsync_NC{args.num_clients}_lr{args.lr_type}_prob{orig_prob}_priori{args.priori}_server{args.server_time}.pkl', 'wb'))


            #print('Iter [%d/%d]: Train Acc: %.2f - Test Acc: %.2f - Loss: %.4f' 
            #        % (step + 1, train_size, avg_train_acc.item(), avg_test_acc.item(), avg_loss.item()))
            print('Iter [%d/%d]: Train Acc: %.2f - Loss: %.4f' 
                    % (step + 1, train_size, avg_train_acc.item(), avg_loss.item()))
            import pdb; pdb.set_trace()

    print('Time taken: %.2f sec.' % (time.time() - start))

