import numpy as np
import pandas as pd
import sys
import tensorflow as tf
from tensorflow.keras import layers

X_train = pd.read_csv('train.csv')
X_test = pd.read_csv('test.csv')
X_valid = pd.read_csv('valid.csv')

Y_train = X_train.pop('stabf')
Y_train = pd.get_dummies(Y_train)

Y_test = X_test.pop('stabf')
Y_test = pd.get_dummies(Y_test)

Y_valid = X_valid.pop('stabf')
Y_valid = pd.get_dummies(Y_valid)

model = tf.keras.Sequential([
                          layers.Input(shape=(12,)),
                          layers.Dense(32),
                          layers.Dense(16),
                          layers.Dense(2, activation='softmax')
])

model.compile(
    loss=tf.keras.losses.BinaryCrossentropy(),
    optimizer=tf.keras.optimizers.Adam(lr=float(sys.argv[1])),
    metrics=[tf.keras.metrics.BinaryAccuracy()])


history = model.fit(tf.convert_to_tensor(X_train, np.float32),
                    Y_train, epochs=2, validation_data=(X_valid, Y_valid))

model.save('grid-stability-dense.h5')