# train linear regression model on random data and test it securely

program.options_from_args()

from sklearn.datasets import make_regression
from sklearn.model_selection import train_test_split

X, y = make_regression(n_samples=1000, n_targets=2)

if len(y.shape) == 1:
    import numpy
    y = numpy.asmatrix(y).transpose()

X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=0)

import torch
import torch.nn as nn

net = nn.Sequential(
    nn.Linear(X.shape[1], y_test.shape[1])
)

# train for a bit
optimizer = torch.optim.SGD(net.parameters(), lr=.1)
criterion = nn.MSELoss()

for i in range(50):
    inputs = torch.Tensor(X_train)
    labels = torch.Tensor(y_train)
    optimizer.zero_grad()
    outputs = net(inputs)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()
    print('Training loss: %f' % loss)

    with torch.no_grad():
        inputs = torch.Tensor(X_test)
        labels = torch.Tensor(y_test)
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        print('Test loss: %f' % loss)

from Compiler import ml

X_train = sfix.input_tensor_via(0, X_train)
X_test = sfix.input_tensor_via(0, X_test)

y_train = sfix.input_tensor_via(0, y_train)
y_test = sfix.input_tensor_via(0, y_test)

layers = ml.layers_from_torch(net, X_train.shape, input_via=0, batch_size=100,
                              regression=True)

ml.set_n_threads(8)

optimizer = ml.Optimizer(layers)

pred = optimizer.eval(X_test)
print_ln('Truth %s', (y_test).reveal_nested())
print_ln('Prediction %s', (pred).reveal_nested())
print_ln('Difference %s', (Matrix.create_from(pred) - y_test).reveal_nested())
print_ln('Secure test loss: %s',
         ((sum((pred[:] - y_test[:]) ** 2)) /
          (y.shape[1] * len(y_test))).reveal())
