from datetime import datetime
import mlflow
from mlflow.models.signature import infer_signature
import pandas as pd
from sacred import Experiment
from sacred.observers import FileStorageObserver, 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):
    with mlflow.start_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(learning_rate=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')
        mlflow.keras.save_model(keras_model=model,
                                path='grid-stability-dense',
                                signature=infer_signature(X_train, Y_train),
                                input_example=X_train.iloc[0])

        _run.info['history'] = str(history.history['loss'][-1])
        mlflow.log_metric('loss', history.history['loss'][-1])
        mlflow.log_param('learning_rate', learning_rate)


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


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