# this trains an AlexNet-like network on CIFAR-10 in cleartext
# before testing it in secure computation

program.options_from_args()

from Compiler import ml

try:
    ml.set_n_threads(int(program.args[1]))
except:
    pass

import torchvision
import torch
import numpy

get_data = lambda train, transform=None: torchvision.datasets.CIFAR10(
    root='/tmp', train=train, download=True, transform=transform)

ds = get_data(False)

# get 100 random samples
indices = numpy.random.randint(len(ds.data), size=(100,))
# normalize to [-1,1] before input
test_samples = sfix.input_tensor_via(
    0, numpy.take(ds.data / 255 * 2 - 1, indices, 0))
test_labels = sint.input_tensor_via(
    0, numpy.take(ds.targets, indices, 0), one_hot=True)

import torch.nn as nn

net = nn.Sequential(
    nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=2),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2),
    nn.Conv2d(64, 96, kernel_size=3, padding=2),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2),
    nn.Conv2d(96, 96, kernel_size=3, padding=1),
    nn.ReLU(),
    nn.Conv2d(96, 64, kernel_size=3, padding=1),
    nn.ReLU(),
    nn.Conv2d(64, 64, kernel_size=3, padding=1),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=3, stride=2),
    nn.Flatten(),
    nn.Linear(1024, 128),
    nn.ReLU(),
    nn.Linear(128, 256),
    nn.ReLU(),
    nn.Linear(256, 10),
)

# train for a bit
transform = torchvision.transforms.Compose(
    [torchvision.transforms.ToTensor(), lambda x: 2 * x - 1])
ds = get_data(train=True, transform=transform)
optimizer = torch.optim.Adam(net.parameters(), amsgrad=True)
criterion = nn.CrossEntropyLoss()

for i, data in enumerate(torch.utils.data.DataLoader(ds, batch_size=128)):
    inputs, labels = data
    optimizer.zero_grad()
    outputs = net(inputs)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()

with torch.no_grad():
    ds = get_data(False, transform)
    total = correct_classified = 0
    for data in torch.utils.data.DataLoader(ds, batch_size=128):
        inputs, labels = data
        outputs = net(inputs)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct_classified += (predicted == labels).sum().item()
    test_acc = (100 * correct_classified / total)
    print('Cleartext test accuracy of the network: %.2f %%' % test_acc)

from Compiler import ml

layers = ml.layers_from_torch(net, test_samples.shape, 128, input_via=0)

optimizer = ml.SGD(layers)

# output to be used in Scripts/torch_cifar_alex_import.py
optimizer.reveal_model_to_binary()

n_correct, loss = optimizer.reveal_correctness(test_samples, test_labels, 128)
print_ln('Secure accuracy: %s (%s/%s)', cfix(n_correct) / len(test_samples),
         n_correct, len(test_samples))
