from sklearn import datasets, linear_model
from sklearn.model_selection import train_test_split

X, y = datasets.load_diabetes(return_X_y=True)

# normalize
y /= y.max()

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

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

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

from Compiler import ml

try:
    batch_size = int(program.args[1])
except:
    batch_size = 1

linear = ml.SGDLinear(100, batch_size, program)

linear.fit(X_train, y_train)
print_ln('model %s', linear.opt.layers[0].W[:].reveal())
print_ln('diff %s', (linear.predict(X_test) - y_test).reveal())

linear.fit_with_testing(X_train, y_train, X_test, y_test)
print_ln('model %s', linear.opt.layers[0].W[:].reveal())
print_ln('diff %s', (linear.predict(X_test) - y_test).reveal())
