# this trains LeNet on CIFAR-10

program.options_from_args()

from Compiler import ml

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

import torchvision, numpy
data = []
for train in True, False:
    ds = torchvision.datasets.CIFAR10(root='/tmp', train=train, download=True)
    # normalize to [-1,1] before input
    samples = sfix.input_tensor_via(0, ds.data / 255 * 2 - 1, binary=True)
    labels = sint.input_tensor_via(0, ds.targets, binary=True, one_hot=True)
    data += [(labels, samples)]

(training_labels, training_samples), (test_labels, test_samples) = data

import torch
import torch.nn as nn

net = nn.Sequential(
    nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=2),
    nn.ReLU(),
    nn.BatchNorm2d(64),
    nn.MaxPool2d(kernel_size=2),
    nn.Conv2d(64, 96, kernel_size=3, padding=2),
    nn.ReLU(),
    nn.BatchNorm2d(96),
    nn.MaxPool2d(kernel_size=2),
    nn.Conv2d(96, 96, kernel_size=3, padding=1),
    nn.ReLU(),
    nn.BatchNorm2d(96),
    nn.Conv2d(96, 64, kernel_size=3, padding=1),
    nn.ReLU(),
    nn.BatchNorm2d(64),
    nn.Conv2d(64, 64, kernel_size=3, padding=1),
    nn.ReLU(),
    nn.BatchNorm2d(64),
    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),
)

# test network
ds = torchvision.datasets.CIFAR10(
    root='/tmp', transform=torchvision.transforms.ToTensor())
inputs = next(iter(torch.utils.data.DataLoader(ds)))[0]
print(inputs.shape)
outputs = net(inputs)

layers = ml.layers_from_torch(net, training_samples.shape, 128)

optimizer = ml.SGD(layers)
optimizer.fit(
    training_samples,
    training_labels,
    epochs=int(program.args[1]),
    batch_size=128,
    validation_data=(test_samples, test_labels),
    program=program
)
