from datetime import datetime
import pandas as pd
from sacred import Experiment
from sacred.observers import MongoObserver
import sys
import tensorflow
from tensorflow.keras import layers

ex = Experiment("470607", interactive=False, save_git_info=False)
ex.observers.append(MongoObserver(url='mongodb://mongo_user:mongo_password_IUM_2021@172.17.0.1:27017', db_name='sacred'))
ex.observers.append(FileStorageObserver('my_runs'))

@ex.config
def my_config():
    learning_rate = float(sys.argv[1])

@ex.capture
def prepare_train_model(learning_rate, _run):
    _run.info["prepare_model"] = str(datetime.now())

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

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

    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=learning_rate),
        metrics=[tensorflow.keras.metrics.BinaryAccuracy()])

    history = model.fit(X_train, Y_train, epochs=2, validation_data=(X_valid, Y_valid))

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

    _run['history'] = history

@ex.main
def my_main(learning_rate):
    print(prepare_train_model())


r = ex.run()
ex.add_artifact('grid-stability-dense.h5')