import csv
import pandas as pd
import seaborn as sns
import sys
import tensorflow
from tensorflow.keras import layers
# from tensorflow.keras.models import load_model


# X_test = pd.read_csv('test.csv')
#
# Y_test = X_test.pop('stabf')
# Y_test = pd.get_dummies(Y_test)
#
# model = load_model('grid-stability-dense.h5')
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 = tensorflow.keras.Sequential([
    layers.Input(shape=(12,)),
    layers.Dense(32),
    layers.Dense(16),
    layers.Dense(2, activation='softmax')
])

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

history = model.fit(X_train, Y_train, epochs=2, validation_data=(X_valid, Y_valid))
results = model.evaluate(X_test, Y_test, batch_size=64)

with open('eval.csv', 'a', newline='') as fp:
    wr = csv.writer(fp, dialect='excel')
    wr.writerow(results)

sns.set_theme(style="darkgrid")
df = pd.read_csv('eval.csv')
sns.lineplot(x='build', y='score', data=df.iloc[1])