import numpy as np
import pandas as pd
import tensorflow as tf

from sklearn import preprocessing
from sklearn.model_selection import train_test_split
from tensorflow.keras import layers


def onezero(label):
  return 0 if label == 'unstable' else 1


df = pd.read_csv('smart_grid_stability_augmented.csv')

scaler = preprocessing.StandardScaler().fit(df.iloc[:, 0:-1])
df_norm_array = scaler.transform(df.iloc[:, 0:-1])
df_norm = pd.DataFrame(data=df_norm_array,
                       columns=df.columns[:-1])
df_norm['stabf'] = df['stabf']

df_norm_data = df_norm.copy()
df_norm_data = df_norm_data.drop('stab', axis=1)
df_norm_labels = df_norm_data.pop('stabf')

X_train, X_testAndValid, Y_train, Y_testAndValid = train_test_split(
    df_norm_data,
    df_norm_labels,
    test_size=0.2,
    random_state=42)

X_test, X_valid, Y_test, Y_valid = train_test_split(
    X_testAndValid,
    Y_testAndValid,
    test_size=0.5,
    random_state=42)

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

model.compile(
    loss=tf.losses.BinaryCrossentropy(),
    optimizer=tf.optimizers.Adam(),
    metrics=[tf.keras.metrics.BinaryAccuracy()])

Y_train_one_zero = [onezero(x) for x in Y_train]
Y_train_onehot = np.eye(2)[Y_train_one_zero]

Y_test_one_zero = [onezero(x) for x in Y_test]
Y_test_onehot = np.eye(2)[Y_test_one_zero]

history = model.fit(tf.convert_to_tensor(X_train, np.float32), Y_train_onehot, epochs=5)

results = model.evaluate(X_test, Y_test_onehot, batch_size=64)
f = open('model_eval.txt', 'w')
f.write('test loss: ' + str(results[0]) + '\n' + 'test acc: ' + str(results[1]))
f.close()