"""
Train 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
import torchvision.transforms as transforms

import argparse
import numpy as np
import time
import os
import copy
import random
import pickle
import math
import itertools

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_small import *
from models.mvcnn_bottom_small import *
from models.resnet import *
from models.resnet_top import *
import util
from logger import Logger
from custom_dataset import MultiViewDataSet
import sys

from sklearn.cluster import KMeans
from sklearn import metrics as skmetrics

import latbin
from PIL import Image

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

# 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=8, type=int,
                    metavar='N', help='mini-batch size (default: 8)')
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('--local_epochs', type=int, help='Number of local epochs to run at each client before synchronizing',
                    default=1)
parser.add_argument('--labeled_frac', type=float, help='Fraction of full dataset that is labeled.',
                    default=0.25)
parser.add_argument('--weight_decay', type=float, help='Fraction of full dataset that is labeled.',
                    default=1)
parser.add_argument('--attack', type=str, help='Information available for the attack',
                    default="none")
parser.add_argument('--dataset', type=int, help='Dataset to use',
                    default=0)
parser.add_argument('--seed', type=int, help='Random seed to use', default=42)

# Parse input arguments
args = parser.parse_args()

np.random.seed(args.seed)
torch.manual_seed(args.seed)
random.seed(args.seed)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

print('Loading data')

if args.dataset == 0:
    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)*args.labeled_frac)])
    samples_weight = torch.tensor([0.1, 0.5, 0.1, 0.1, 0.7, 0.1, 0.1, 0.1, 0.1, 1.0])
    num_iters = math.ceil(len(dset_train_sub)/args.batch_size)
    sampler = torch.utils.data.WeightedRandomSampler(samples_weight, num_iters*args.batch_size)
    train_loader = DataLoader(dset_train_sub, batch_size=args.batch_size, shuffle=False, num_workers=1, sampler=sampler)

    dset_aux = dset_train #torch.utils.data.Subset(dset_train, indices[int(len(dset_train)*args.labeled_frac):int(2*len(dset_train)*args.labeled_frac)])

    dset_val = MultiViewDataSet(args.data, 'test', transform=transform)
    test_loader = DataLoader(dset_val, batch_size=args.batch_size, shuffle=False, num_workers=1)
    classes = dset_train.classes
else:
    transform = transforms.Compose(
    [transforms.ToTensor()])#,transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

    dset_train = torchvision.datasets.CIFAR100(root='./data', train=True,
                                                    download=True, transform=transform)
    indices = torch.randperm(len(dset_train))
    dset_train_sub = torch.utils.data.Subset(dset_train, indices[:int(len(dset_train)*0.05)])#*args.labeled_frac)])
    samples_weight = torch.tensor([0.1, 0.5, 0.1, 0.1, 0.7, 0.1, 0.1, 0.1, 0.1, 1.0])
    samples_weight = samples_weight.repeat(10)
    num_iters = math.ceil(len(dset_train_sub)/args.batch_size)
    sampler = torch.utils.data.WeightedRandomSampler(samples_weight, num_iters*args.batch_size)
    train_loader = DataLoader(dset_train_sub, batch_size=args.batch_size, shuffle=False, num_workers=1, sampler=sampler)

    indices = torch.randperm(len(dset_train))
    dset_aux = dset_train #torch.utils.data.Subset(dset_train, indices[int(len(dset_train)*args.labeled_frac):int(2*len(dset_train)*args.labeled_frac)])

    dset_val = torchvision.datasets.CIFAR100(root='./data', train=False,
                                                    download=True, transform=transform)
    test_loader = DataLoader(dset_val, batch_size=args.batch_size, shuffle=False, num_workers=1)

    classes = dset_train.classes 

num_classes = len(classes)
print(len(classes), classes)

# Loss and Optimizer
n_epochs = args.epochs
criterion = nn.CrossEntropyLoss()
if args.dataset == 0:
    coords_per = int(12/num_clients)
else:
    coords_per = int(32*32/num_clients)

best_acc = 0.0
best_loss = 0.0
start_epoch = 0

losses = []
accs_train = []
accs_test = []

models = []
optimizers = []
# Make models for each client
for i in range(num_clients+1):
    if i == num_clients:
        if args.dataset == 0:
            model = mvcnn_top(pretrained=args.pretrained, 
                            num_classes=len(classes), num_clients=num_clients)
        else:
            model = resnet_top(pretrained=args.pretrained, 
                            num_classes=len(classes), num_clients=num_clients)
    else:
        if args.dataset == 0:
            model = mvcnn_bottom(pretrained=args.pretrained,num_classes=len(classes))
        else:
            #model = resnet18(pretrained=args.pretrained, num_classes=len(classes))
            model = torch.hub.load('pytorch/vision:v0.5.0', 'resnet18', pretrained=False, num_classes=len(classes))

    model.to(device)
    cudnn.benchmark = True

    #optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=0.1)

    # Check if model .pt exists
    #PATH = f"/gpfs/u/home/VFLA/VFLAcstg/scratch/checkpoint{i}_supervised_NC{args.num_clients}_frac{args.labeled_frac}_seed{args.seed}.pt"
    #if os.path.exists(PATH):
    #    print("Loading from checkpoint")
    #    checkpoint = torch.load(PATH) 
    #    model.load_state_dict(checkpoint['model_state_dict'])
    #    model.eval()
    #    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    #    start_epoch = checkpoint['epoch']
    #    losses = pickle.load(open(f'/gpfs/u/home/VFLA/VFLAcstg/scratch/loss_mvcnn_NC{args.num_clients}_LE{args.local_epochs}_frac{args.labeled_frac}_seed{args.seed}.pkl','rb'))
    #    accs_train = pickle.load(open(f'/gpfs/u/home/VFLA/VFLAcstg/scratch/accs_train_mvcnn_NC{args.num_clients}_LE{args.local_epochs}_frac{args.labeled_frac}_seed{args.seed}.pkl','rb'))
    #    accs_test = pickle.load(open(f'/gpfs/u/home/VFLA/VFLAcstg/scratch/accs_test_mvcnn_NC{args.num_clients}_LE{args.local_epochs}_frac{args.labeled_frac}_seed{args.seed}.pkl','rb'))

    models.append(model)
    optimizers.append(optimizer)

def save_eval(models, train_loader, test_loader, losses, accs_train, accs_test, step, train_size, leakages):
    """
    Evaluate and save current loss and 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_mvcnn_NC{args.num_clients}_LE{args.local_epochs}_frac{args.labeled_frac}_batch{args.batch_size}_dataset{args.dataset}_seed{args.seed}.pkl', 'wb'))
    pickle.dump(accs_train, open(f'accs_train_mvcnn_NC{args.num_clients}_LE{args.local_epochs}_frac{args.labeled_frac}_batch{args.batch_size}_dataset{args.dataset}_seed{args.seed}.pkl', 'wb'))
    pickle.dump(accs_test, open(f'accs_test_mvcnn_NC{args.num_clients}_LE{args.local_epochs}_frac{args.labeled_frac}_batch{args.batch_size}_dataset{args.dataset}_seed{args.seed}.pkl', 'wb'))
    pickle.dump(leakages, open(f'leakage_mvcnn_NC{args.num_clients}_LE{args.local_epochs}_frac{args.labeled_frac}_batch{args.batch_size}_dataset{args.dataset}_seed{args.seed}.pkl', 'wb'))

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

def labels_estimate(models, optimizers):
    E = []
    S = np.zeros(num_classes)
    G = np.zeros(num_classes) 
    m = 0 
    grads = []
    for param in models[-1].parameters():
        grads.append(param.grad)
    for i in range(num_classes):
        G[i] = torch.sum(grads[0][i,:])
    if args.attack == "white-box" or args.attack == "aux-info":
        # Load aux data 
        Gbar = np.zeros(num_classes) 
        for clas in range(num_classes):
            if args.dataset == 0:
                dset_aux_x = np.array(dset_aux.x)
                dset_aux_y = dset_aux.y
            else:
                dset_aux_x = dset_aux.data
                dset_aux_y = dset_aux.targets
            index = np.where(np.array(dset_aux_y) == clas)
            aux_x = (dset_aux_x[index])[:args.batch_size]
            if args.dataset == 0:
                aux_x_new = []
                for original_views in aux_x:
                    views = []
                    for view in original_views:
                        im = Image.open(view)
                        im = im.convert('RGB')
                        if transform is not None:
                            im = transform(im)
                        views.append(im)
                    aux_x_new.append(torch.stack(views))
                aux_x = np.array(torch.stack(aux_x_new))
            aux_y = torch.tensor((np.array(dset_aux_y)[index])[:args.batch_size])
        
            # Feed aux data into white-box model
            forward(models, optimizers, aux_x, aux_y)
            grads = []
            for param in models[-1].parameters():
                grads.append(param.grad)
            Gbar[clas] = torch.sum(grads[0][clas,:])

        # Estimate m
        for j, gbar_i in enumerate(Gbar):
            m += (1/(args.batch_size*num_classes))*gbar_i*(1+1/num_classes) 

        t = 10
        Gbar = np.zeros(num_classes) 
        for i in range(num_classes):
            for k in range(t):
                ik = random.choice([num for num in range(0,9) if num != i]) 
                if args.dataset == 0:
                    dset_aux_x = np.array(dset_aux.x)
                    dset_aux_y = dset_aux.y
                else:
                    dset_aux_x = dset_aux.data
                    dset_aux_y = dset_aux.targets
                index = np.where(np.array(dset_aux_y) == ik)
                aux_x = (dset_aux_x[index])[:args.batch_size]
                if args.dataset == 0:
                    aux_x_new = []
                    for original_views in aux_x:
                        views = []
                        for view in original_views:
                            im = Image.open(view)
                            im = im.convert('RGB')
                            if transform is not None:
                                im = transform(im)
                            views.append(im)
                        aux_x_new.append(torch.stack(views))
                    aux_x = np.array(torch.stack(aux_x_new))
                aux_y = torch.tensor((np.array(dset_aux_y)[index])[:args.batch_size])

                # Feed aux data into white-box model
                forward(models, optimizers, aux_x, aux_y)
                grads = []
                for param in models[-1].parameters():
                    grads.append(param.grad)
                Gbar[i] += torch.sum(grads[0][i,:])

        # Estimate s
        for j, gbar_i in enumerate(Gbar):
            S[j] = (1/t)*gbar_i 

    elif args.attack != "shared-only":
        for i, g_i in enumerate(G):
            if g_i < 0:
                m += (1/args.batch_size)*g_i*(1+1/num_classes) 
    if args.attack != "none":
        for i, g_i in enumerate(G):
            if g_i < 0:
                E.append(i)
                g_i -= m
        G = G - S
        while len(E) < args.batch_size:
            i = np.argmin(G)
            E.append(i)
            G[i] -= m

    return E


def leakage_percent(E, y):
    num_correct = 0
    y = copy.deepcopy(y.cpu().detach().numpy())
    num_labels = len(y)
    for guess in E:
        if guess in y:
            y = np.delete(y, np.argwhere(y == guess)[0])
            num_correct += 1
    return num_correct/num_labels


def forward(models, optimizers, inputs, targets):
    """
    Train all clients on all batches 
    """

    server_model = models[-1]
    server_optimizer = optimizers[-1]

    #inputs = np.stack(inputs, axis=1)
    inputs = torch.from_numpy(inputs).type(torch.FloatTensor)
    inputs, targets = inputs.cuda(device), targets.cuda(device)
    inputs, targets = Variable(inputs), Variable(targets)

    # Exchange embeddings
    H_orig = [None] * num_clients
    for i in range(num_clients):
        if args.dataset == 0:
            x_local = inputs[:,coords_per*i:coords_per*(i+1),:,:,:]
        else:
            r = math.floor(i/2)
            c = i % 2
            section = int(math.sqrt(coords_per))
            x_local = inputs[:,
                            section*r:section*(r+1),
                            section*c:section*(c+1),:]
            x_local = torch.transpose(x_local,1,3)
        H_orig[i] = models[i](x_local)

    # Train clients
    for i in range(num_clients):
        if args.dataset == 0:
            x_local = inputs[:,coords_per*i:coords_per*(i+1),:,:,:]
        else:
            r = math.floor(i/2)
            c = i % 2
            section = int(math.sqrt(coords_per))
            x_local = inputs[:,
                            section*r:section*(r+1),
                            section*c:section*(c+1),:]
            x_local = torch.transpose(x_local,1,3)
        H = H_orig.copy()
        model = models[i]
        optimizer = optimizers[i]

        # Calculate number of local iterations
        client_epochs = args.local_epochs
        # Train
        # compute output
        outputs = model(x_local)
        H[i] = outputs
        outputs = server_model(torch.cat(H,axis=1))
        loss = criterion(outputs, targets)

        # compute gradient and do gradient step
        optimizer.zero_grad()
        server_optimizer.zero_grad()
        loss.backward(retain_graph=True)
        optimizer.step()

    # Train server
    H = H_orig.copy()
    # compute output
    outputs = server_model(torch.cat(H,axis=1))
    loss = criterion(outputs, targets)

    # compute gradient and do SGD step
    server_optimizer.zero_grad()
    loss.backward(retain_graph=True)

def train(models, optimizers, epoch, leakages): #, centers):
    """
    Train all clients on all batches 
    """

    train_size = len(train_loader)
    server_model = models[-1]
    server_optimizer = optimizers[-1]

    Hs = np.empty((len(train_loader), num_clients), dtype=object)
    Hs.fill([])

    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)
        # Exchange embeddings
        H_orig = [None] * num_clients
        for i in range(num_clients):
            if args.dataset == 0:
                x_local = inputs[:,coords_per*i:coords_per*(i+1),:,:,:]
            else:
                r = math.floor(i/2)
                c = i % 2
                section = int(math.sqrt(coords_per))
                x_local = inputs[:,:,
                                section*r:section*(r+1),
                                section*c:section*(c+1)]
                x_local = torch.transpose(x_local,0,1)
            H_orig[i] = models[i](x_local)
            Hs[step, i] = H_orig[i].cpu().detach().numpy()

        # Train clients
        for i in range(num_clients):
            if args.dataset == 0:
                x_local = inputs[:,coords_per*i:coords_per*(i+1),:,:,:]
            else:
                r = math.floor(i/2)
                c = i % 2
                section = int(math.sqrt(coords_per))
                x_local = inputs[:,:,
                                section*r:section*(r+1),
                                section*c:section*(c+1)]
                x_local = torch.transpose(x_local,0,1)
            H = H_orig.copy()
            model = models[i]
            optimizer = optimizers[i]

            # Calculate number of local iterations
            client_epochs = args.local_epochs
            # Train
            for le in range(client_epochs):
                # compute output
                outputs = model(x_local)
                H[i] = outputs
                outputs = server_model(torch.cat(H,axis=1))
                loss = criterion(outputs, targets)

                # compute gradient and do gradient step
                optimizer.zero_grad()
                server_optimizer.zero_grad()
                loss.backward(retain_graph=True)
                optimizer.step()

        # Train server
        for le in range(args.local_epochs):
            H = H_orig.copy()
            # compute output
            outputs = server_model(torch.cat(H,axis=1))
            loss = criterion(outputs, targets)

            # compute gradient and do SGD step
            server_optimizer.zero_grad()
            loss.backward(retain_graph=True)
            server_optimizer.step()

            # Estimate labels based on final layer gradients
            E = labels_estimate(models, optimizers)
            percent_correct = leakage_percent(E, targets)
            leakages.append([])
            leakages[-1].append(percent_correct)
            percent_correct = leakage_percent(np.random.randint(0,9,size=targets.shape), targets)
            leakages[-1].append(percent_correct)

        if (step + 1) % args.print_freq == 0:
            print("\tServer Iter [%d/%d] Loss: %.4f - Leakage: %.1f - Guess: %.1f" % (step + 1, train_size, loss.item(), leakages[-1][0]*100, leakages[-1][1]*100))

    return leakages


# 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):
                if args.dataset == 0:
                    x_local = inputs[:,coords_per*i:coords_per*(i+1),:,:,:]
                else:
                    r = math.floor(i/2)
                    c = i % 2
                    section = int(math.sqrt(coords_per))
                    x_local = inputs[:,:,
                                    section*r:section*(r+1),
                                    section*c:section*(c+1)]
                    x_local = torch.transpose(x_local,0,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

# Get initial loss/accuracy
#if start_epoch == 0:
#    save_eval(models, train_loader, test_loader, losses, accs_train, accs_test, 0, len(train_loader))
# Training / Eval loop
leakages = []
train_size = len(train_loader)
for epoch in range(start_epoch, n_epochs):
    print('\n-----------------------------------')
    print('Epoch: [%d/%d]' % (epoch+1, n_epochs))
    start = time.time()

    leakages = train(models, optimizers, epoch, leakages)
    save_eval(models, train_loader, test_loader, losses, accs_train, accs_test, epoch, train_size, leakages)
    #for i in range(num_clients+1):
    #    PATH = f"/gpfs/u/home/VFLA/VFLAcstg/scratch/checkpoint{i}_supervised_NC{args.num_clients}_frac{args.labeled_frac}_seed{args.seed}.pt"

    #    torch.save({
    #                'epoch': epoch+1,
    #                'model_state_dict': models[i].state_dict(),
    #                'optimizer_state_dict': optimizers[i].state_dict(),
    #                'loss': 0,
    #                }, PATH)
    print('Time taken: %.2f sec.' % (time.time() - start))


"""
Train 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
import torchvision.transforms as transforms

import argparse
import numpy as np
import time
import os
import copy
import random
import pickle
import math
import itertools

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_small import *
from models.mvcnn_bottom_small import *
from models.resnet import *
from models.resnet_top import *
import util
from logger import Logger
from custom_dataset import MultiViewDataSet
import sys

from sklearn.cluster import KMeans
from sklearn import metrics as skmetrics

import latbin
from PIL import Image

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

# 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=8, type=int,
                    metavar='N', help='mini-batch size (default: 8)')
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('--local_epochs', type=int, help='Number of local epochs to run at each client before synchronizing',
                    default=1)
parser.add_argument('--labeled_frac', type=float, help='Fraction of full dataset that is labeled.',
                    default=0.25)
parser.add_argument('--weight_decay', type=float, help='Fraction of full dataset that is labeled.',
                    default=1)
parser.add_argument('--attack', type=str, help='Information available for the attack',
                    default="none")
parser.add_argument('--dataset', type=int, help='Dataset to use',
                    default=0)
parser.add_argument('--seed', type=int, help='Random seed to use', default=42)

# Parse input arguments
args = parser.parse_args()

np.random.seed(args.seed)
torch.manual_seed(args.seed)
random.seed(args.seed)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

print('Loading data')

if args.dataset == 0:
    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)*args.labeled_frac)])
    samples_weight = torch.tensor([0.1, 0.5, 0.1, 0.1, 0.7, 0.1, 0.1, 0.1, 0.1, 1.0])
    num_iters = math.ceil(len(dset_train_sub)/args.batch_size)
    sampler = torch.utils.data.WeightedRandomSampler(samples_weight, num_iters*args.batch_size)
    train_loader = DataLoader(dset_train_sub, batch_size=args.batch_size, shuffle=False, num_workers=1, sampler=sampler)

    dset_aux = dset_train #torch.utils.data.Subset(dset_train, indices[int(len(dset_train)*args.labeled_frac):int(2*len(dset_train)*args.labeled_frac)])

    dset_val = MultiViewDataSet(args.data, 'test', transform=transform)
    test_loader = DataLoader(dset_val, batch_size=args.batch_size, shuffle=False, num_workers=1)
    classes = dset_train.classes
else:
    transform = transforms.Compose(
    [transforms.ToTensor()])#,transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

    dset_train = torchvision.datasets.CIFAR100(root='./data', train=True,
                                                    download=True, transform=transform)
    indices = torch.randperm(len(dset_train))
    dset_train_sub = torch.utils.data.Subset(dset_train, indices[:int(len(dset_train)*0.05)])#*args.labeled_frac)])
    samples_weight = torch.tensor([0.1, 0.5, 0.1, 0.1, 0.7, 0.1, 0.1, 0.1, 0.1, 1.0])
    samples_weight = samples_weight.repeat(10)
    num_iters = math.ceil(len(dset_train_sub)/args.batch_size)
    sampler = torch.utils.data.WeightedRandomSampler(samples_weight, num_iters*args.batch_size)
    train_loader = DataLoader(dset_train_sub, batch_size=args.batch_size, shuffle=False, num_workers=1, sampler=sampler)

    indices = torch.randperm(len(dset_train))
    dset_aux = dset_train #torch.utils.data.Subset(dset_train, indices[int(len(dset_train)*args.labeled_frac):int(2*len(dset_train)*args.labeled_frac)])

    dset_val = torchvision.datasets.CIFAR100(root='./data', train=False,
                                                    download=True, transform=transform)
    test_loader = DataLoader(dset_val, batch_size=args.batch_size, shuffle=False, num_workers=1)

    classes = dset_train.classes 

num_classes = len(classes)
print(len(classes), classes)

# Loss and Optimizer
n_epochs = args.epochs
criterion = nn.CrossEntropyLoss()
if args.dataset == 0:
    coords_per = int(12/num_clients)
else:
    coords_per = int(32*32/num_clients)

best_acc = 0.0
best_loss = 0.0
start_epoch = 0

losses = []
accs_train = []
accs_test = []

models = []
optimizers = []
# Make models for each client
for i in range(num_clients+1):
    if i == num_clients:
        if args.dataset == 0:
            model = mvcnn_top(pretrained=args.pretrained, 
                            num_classes=len(classes), num_clients=num_clients)
        else:
            model = resnet_top(pretrained=args.pretrained, 
                            num_classes=len(classes), num_clients=num_clients)
    else:
        if args.dataset == 0:
            model = mvcnn_bottom(pretrained=args.pretrained,num_classes=len(classes))
        else:
            #model = resnet18(pretrained=args.pretrained, num_classes=len(classes))
            model = torch.hub.load('pytorch/vision:v0.5.0', 'resnet18', pretrained=False, num_classes=len(classes))

    model.to(device)
    cudnn.benchmark = True

    #optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=0.1)

    # Check if model .pt exists
    #PATH = f"/gpfs/u/home/VFLA/VFLAcstg/scratch/checkpoint{i}_supervised_NC{args.num_clients}_frac{args.labeled_frac}_seed{args.seed}.pt"
    #if os.path.exists(PATH):
    #    print("Loading from checkpoint")
    #    checkpoint = torch.load(PATH) 
    #    model.load_state_dict(checkpoint['model_state_dict'])
    #    model.eval()
    #    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    #    start_epoch = checkpoint['epoch']
    #    losses = pickle.load(open(f'/gpfs/u/home/VFLA/VFLAcstg/scratch/loss_mvcnn_NC{args.num_clients}_LE{args.local_epochs}_frac{args.labeled_frac}_seed{args.seed}.pkl','rb'))
    #    accs_train = pickle.load(open(f'/gpfs/u/home/VFLA/VFLAcstg/scratch/accs_train_mvcnn_NC{args.num_clients}_LE{args.local_epochs}_frac{args.labeled_frac}_seed{args.seed}.pkl','rb'))
    #    accs_test = pickle.load(open(f'/gpfs/u/home/VFLA/VFLAcstg/scratch/accs_test_mvcnn_NC{args.num_clients}_LE{args.local_epochs}_frac{args.labeled_frac}_seed{args.seed}.pkl','rb'))

    models.append(model)
    optimizers.append(optimizer)

def save_eval(models, train_loader, test_loader, losses, accs_train, accs_test, step, train_size, leakages):
    """
    Evaluate and save current loss and 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_mvcnn_NC{args.num_clients}_LE{args.local_epochs}_frac{args.labeled_frac}_batch{args.batch_size}_dataset{args.dataset}_seed{args.seed}.pkl', 'wb'))
    pickle.dump(accs_train, open(f'accs_train_mvcnn_NC{args.num_clients}_LE{args.local_epochs}_frac{args.labeled_frac}_batch{args.batch_size}_dataset{args.dataset}_seed{args.seed}.pkl', 'wb'))
    pickle.dump(accs_test, open(f'accs_test_mvcnn_NC{args.num_clients}_LE{args.local_epochs}_frac{args.labeled_frac}_batch{args.batch_size}_dataset{args.dataset}_seed{args.seed}.pkl', 'wb'))
    pickle.dump(leakages, open(f'leakage_mvcnn_NC{args.num_clients}_LE{args.local_epochs}_frac{args.labeled_frac}_batch{args.batch_size}_dataset{args.dataset}_seed{args.seed}.pkl', 'wb'))

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

def labels_estimate(models, optimizers):
    E = []
    S = np.zeros(num_classes)
    G = np.zeros(num_classes) 
    m = 0 
    grads = []
    for param in models[-1].parameters():
        grads.append(param.grad)
    for i in range(num_classes):
        G[i] = torch.sum(grads[0][i,:])
    if args.attack == "white-box" or args.attack == "aux-info":
        # Load aux data 
        Gbar = np.zeros(num_classes) 
        for clas in range(num_classes):
            if args.dataset == 0:
                dset_aux_x = np.array(dset_aux.x)
                dset_aux_y = dset_aux.y
            else:
                dset_aux_x = dset_aux.data
                dset_aux_y = dset_aux.targets
            index = np.where(np.array(dset_aux_y) == clas)
            aux_x = (dset_aux_x[index])[:args.batch_size]
            if args.dataset == 0:
                aux_x_new = []
                for original_views in aux_x:
                    views = []
                    for view in original_views:
                        im = Image.open(view)
                        im = im.convert('RGB')
                        if transform is not None:
                            im = transform(im)
                        views.append(im)
                    aux_x_new.append(torch.stack(views))
                aux_x = np.array(torch.stack(aux_x_new))
            aux_y = torch.tensor((np.array(dset_aux_y)[index])[:args.batch_size])
        
            # Feed aux data into white-box model
            forward(models, optimizers, aux_x, aux_y)
            grads = []
            for param in models[-1].parameters():
                grads.append(param.grad)
            Gbar[clas] = torch.sum(grads[0][clas,:])

        # Estimate m
        for j, gbar_i in enumerate(Gbar):
            m += (1/(args.batch_size*num_classes))*gbar_i*(1+1/num_classes) 

        t = 10
        Gbar = np.zeros(num_classes) 
        for i in range(num_classes):
            for k in range(t):
                ik = random.choice([num for num in range(0,9) if num != i]) 
                if args.dataset == 0:
                    dset_aux_x = np.array(dset_aux.x)
                    dset_aux_y = dset_aux.y
                else:
                    dset_aux_x = dset_aux.data
                    dset_aux_y = dset_aux.targets
                index = np.where(np.array(dset_aux_y) == ik)
                aux_x = (dset_aux_x[index])[:args.batch_size]
                if args.dataset == 0:
                    aux_x_new = []
                    for original_views in aux_x:
                        views = []
                        for view in original_views:
                            im = Image.open(view)
                            im = im.convert('RGB')
                            if transform is not None:
                                im = transform(im)
                            views.append(im)
                        aux_x_new.append(torch.stack(views))
                    aux_x = np.array(torch.stack(aux_x_new))
                aux_y = torch.tensor((np.array(dset_aux_y)[index])[:args.batch_size])

                # Feed aux data into white-box model
                forward(models, optimizers, aux_x, aux_y)
                grads = []
                for param in models[-1].parameters():
                    grads.append(param.grad)
                Gbar[i] += torch.sum(grads[0][i,:])

        # Estimate s
        for j, gbar_i in enumerate(Gbar):
            S[j] = (1/t)*gbar_i 

    elif args.attack != "shared-only":
        for i, g_i in enumerate(G):
            if g_i < 0:
                m += (1/args.batch_size)*g_i*(1+1/num_classes) 
    if args.attack != "none":
        for i, g_i in enumerate(G):
            if g_i < 0:
                E.append(i)
                g_i -= m
        G = G - S
        while len(E) < args.batch_size:
            i = np.argmin(G)
            E.append(i)
            G[i] -= m

    return E


def leakage_percent(E, y):
    num_correct = 0
    y = copy.deepcopy(y.cpu().detach().numpy())
    num_labels = len(y)
    for guess in E:
        if guess in y:
            y = np.delete(y, np.argwhere(y == guess)[0])
            num_correct += 1
    return num_correct/num_labels


def forward(models, optimizers, inputs, targets):
    """
    Train all clients on all batches 
    """

    server_model = models[-1]
    server_optimizer = optimizers[-1]

    #inputs = np.stack(inputs, axis=1)
    inputs = torch.from_numpy(inputs).type(torch.FloatTensor)
    inputs, targets = inputs.cuda(device), targets.cuda(device)
    inputs, targets = Variable(inputs), Variable(targets)

    # Exchange embeddings
    H_orig = [None] * num_clients
    for i in range(num_clients):
        if args.dataset == 0:
            x_local = inputs[:,coords_per*i:coords_per*(i+1),:,:,:]
        else:
            r = math.floor(i/2)
            c = i % 2
            section = int(math.sqrt(coords_per))
            x_local = inputs[:,
                            section*r:section*(r+1),
                            section*c:section*(c+1),:]
            x_local = torch.transpose(x_local,1,3)
        H_orig[i] = models[i](x_local)

    # Train clients
    for i in range(num_clients):
        if args.dataset == 0:
            x_local = inputs[:,coords_per*i:coords_per*(i+1),:,:,:]
        else:
            r = math.floor(i/2)
            c = i % 2
            section = int(math.sqrt(coords_per))
            x_local = inputs[:,
                            section*r:section*(r+1),
                            section*c:section*(c+1),:]
            x_local = torch.transpose(x_local,1,3)
        H = H_orig.copy()
        model = models[i]
        optimizer = optimizers[i]

        # Calculate number of local iterations
        client_epochs = args.local_epochs
        # Train
        # compute output
        outputs = model(x_local)
        H[i] = outputs
        outputs = server_model(torch.cat(H,axis=1))
        loss = criterion(outputs, targets)

        # compute gradient and do gradient step
        optimizer.zero_grad()
        server_optimizer.zero_grad()
        loss.backward(retain_graph=True)
        optimizer.step()

    # Train server
    H = H_orig.copy()
    # compute output
    outputs = server_model(torch.cat(H,axis=1))
    loss = criterion(outputs, targets)

    # compute gradient and do SGD step
    server_optimizer.zero_grad()
    loss.backward(retain_graph=True)

def train(models, optimizers, epoch, leakages): #, centers):
    """
    Train all clients on all batches 
    """

    train_size = len(train_loader)
    server_model = models[-1]
    server_optimizer = optimizers[-1]

    Hs = np.empty((len(train_loader), num_clients), dtype=object)
    Hs.fill([])

    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)
        # Exchange embeddings
        H_orig = [None] * num_clients
        for i in range(num_clients):
            if args.dataset == 0:
                x_local = inputs[:,coords_per*i:coords_per*(i+1),:,:,:]
            else:
                r = math.floor(i/2)
                c = i % 2
                section = int(math.sqrt(coords_per))
                x_local = inputs[:,:,
                                section*r:section*(r+1),
                                section*c:section*(c+1)]
                x_local = torch.transpose(x_local,0,1)
            H_orig[i] = models[i](x_local)
            Hs[step, i] = H_orig[i].cpu().detach().numpy()

        # Train clients
        for i in range(num_clients):
            if args.dataset == 0:
                x_local = inputs[:,coords_per*i:coords_per*(i+1),:,:,:]
            else:
                r = math.floor(i/2)
                c = i % 2
                section = int(math.sqrt(coords_per))
                x_local = inputs[:,:,
                                section*r:section*(r+1),
                                section*c:section*(c+1)]
                x_local = torch.transpose(x_local,0,1)
            H = H_orig.copy()
            model = models[i]
            optimizer = optimizers[i]

            # Calculate number of local iterations
            client_epochs = args.local_epochs
            # Train
            for le in range(client_epochs):
                # compute output
                outputs = model(x_local)
                H[i] = outputs
                outputs = server_model(torch.cat(H,axis=1))
                loss = criterion(outputs, targets)

                # compute gradient and do gradient step
                optimizer.zero_grad()
                server_optimizer.zero_grad()
                loss.backward(retain_graph=True)
                optimizer.step()

        # Train server
        for le in range(args.local_epochs):
            H = H_orig.copy()
            # compute output
            outputs = server_model(torch.cat(H,axis=1))
            loss = criterion(outputs, targets)

            # compute gradient and do SGD step
            server_optimizer.zero_grad()
            loss.backward(retain_graph=True)
            server_optimizer.step()

            # Estimate labels based on final layer gradients
            E = labels_estimate(models, optimizers)
            percent_correct = leakage_percent(E, targets)
            leakages.append([])
            leakages[-1].append(percent_correct)
            percent_correct = leakage_percent(np.random.randint(0,9,size=targets.shape), targets)
            leakages[-1].append(percent_correct)

        if (step + 1) % args.print_freq == 0:
            print("\tServer Iter [%d/%d] Loss: %.4f - Leakage: %.1f - Guess: %.1f" % (step + 1, train_size, loss.item(), leakages[-1][0]*100, leakages[-1][1]*100))

    return leakages


# 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):
                if args.dataset == 0:
                    x_local = inputs[:,coords_per*i:coords_per*(i+1),:,:,:]
                else:
                    r = math.floor(i/2)
                    c = i % 2
                    section = int(math.sqrt(coords_per))
                    x_local = inputs[:,:,
                                    section*r:section*(r+1),
                                    section*c:section*(c+1)]
                    x_local = torch.transpose(x_local,0,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

# Get initial loss/accuracy
#if start_epoch == 0:
#    save_eval(models, train_loader, test_loader, losses, accs_train, accs_test, 0, len(train_loader))
# Training / Eval loop
leakages = []
train_size = len(train_loader)
for epoch in range(start_epoch, n_epochs):
    print('\n-----------------------------------')
    print('Epoch: [%d/%d]' % (epoch+1, n_epochs))
    start = time.time()

    leakages = train(models, optimizers, epoch, leakages)
    save_eval(models, train_loader, test_loader, losses, accs_train, accs_test, epoch, train_size, leakages)
    #for i in range(num_clients+1):
    #    PATH = f"/gpfs/u/home/VFLA/VFLAcstg/scratch/checkpoint{i}_supervised_NC{args.num_clients}_frac{args.labeled_frac}_seed{args.seed}.pt"

    #    torch.save({
    #                'epoch': epoch+1,
    #                'model_state_dict': models[i].state_dict(),
    #                'optimizer_state_dict': optimizers[i].state_dict(),
    #                'loss': 0,
    #                }, PATH)
    print('Time taken: %.2f sec.' % (time.time() - start))


Train VFL on ModelNet-10 dataset # Set up input arguments # Parse input arguments # Load dataset #torch.utils.data.Subset(dset_train, indices[int(len(dset_train)*args.labeled_frac):int(2*len(dset_train)*args.labeled_frac)]) #,transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]) #*args.labeled_frac)]) #torch.utils.data.Subset(dset_train, indices[int(len(dset_train)*args.labeled_frac):int(2*len(dset_train)*args.labeled_frac)]) # Loss and Optimizer # Make models for each client #model = resnet18(pretrained=args.pretrained, num_classes=len(classes)) #optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) # Check if model .pt exists #PATH = f"/gpfs/u/home/VFLA/VFLAcstg/scratch/checkpoint{i}_supervised_NC{args.num_clients}_frac{args.labeled_frac}_seed{args.seed}.pt" #if os.path.exists(PATH): #    print("Loading from checkpoint") #    checkpoint = torch.load(PATH) #    model.load_state_dict(checkpoint['model_state_dict']) #    model.eval() #    optimizer.load_state_dict(checkpoint['optimizer_state_dict']) #    start_epoch = checkpoint['epoch'] #    losses = pickle.load(open(f'/gpfs/u/home/VFLA/VFLAcstg/scratch/loss_mvcnn_NC{args.num_clients}_LE{args.local_epochs}_frac{args.labeled_frac}_seed{args.seed}.pkl','rb')) #    accs_train = pickle.load(open(f'/gpfs/u/home/VFLA/VFLAcstg/scratch/accs_train_mvcnn_NC{args.num_clients}_LE{args.local_epochs}_frac{args.labeled_frac}_seed{args.seed}.pkl','rb')) #    accs_test = pickle.load(open(f'/gpfs/u/home/VFLA/VFLAcstg/scratch/accs_test_mvcnn_NC{args.num_clients}_LE{args.local_epochs}_frac{args.labeled_frac}_seed{args.seed}.pkl','rb')) Evaluate and save current loss and accuracy # Load aux data # Feed aux data into white-box model # Estimate m # Feed aux data into white-box model # Estimate s Train all clients on all batches #inputs = np.stack(inputs, axis=1) # Exchange embeddings # Train clients # Calculate number of local iterations # Train # compute output # compute gradient and do gradient step # Train server # compute output # compute gradient and do SGD step #, centers): Train all clients on all batches # Convert from list of 3D to 4D # Exchange embeddings # Train clients # Calculate number of local iterations # Train # compute output # compute gradient and do gradient step # Train server # compute output # compute gradient and do SGD step # Estimate labels based on final layer gradients # Validation and Testing Calculate loss and accuracy for a given data_loader # Convert from list of 3D to 4D # Get current embeddings # compute output # Get initial loss/accuracy #if start_epoch == 0: #    save_eval(models, train_loader, test_loader, losses, accs_train, accs_test, 0, len(train_loader)) # Training / Eval loop #for i in range(num_clients+1): #    PATH = f"/gpfs/u/home/VFLA/VFLAcstg/scratch/checkpoint{i}_supervised_NC{args.num_clients}_frac{args.labeled_frac}_seed{args.seed}.pt" #    torch.save({ #                'epoch': epoch+1, #                'model_state_dict': models[i].state_dict(), #                'optimizer_state_dict': optimizers[i].state_dict(), #                'loss': 0, #                }, PATH)