import pandas
import os
from sacred import Experiment
from sacred.observers import MongoObserver
from sacred.observers import FileStorageObserver
from keras.applications.densenet import layers
import tensorflow

exint = Experiment('z-s434686-training')

exint.observers.append(MongoObserver(url='mongodb://admin:IUM_2021@172.17.0.1:27017',
                                  db_name='sacred'))
exint.observers.append(FileStorageObserver('my_runs'))
@exint.config
def my_config():
    EPOCHS = int(os.environ['EPOCHS'])
    

@exint.main
def my_main(EPOCHS, _run):
    _run.info["epochs"] = EPOCHS,
    normalize = layers.Normalization()
    train_data_x = pandas.read_csv('./X_train.csv')
    _run.info["dataset"] = train_data_x
    games_all = train_data_x.copy()
    games_predict = train_data_x.pop('User_Score')
    normalize.adapt(games_all)
    norm_games_model = tensorflow.keras.Sequential([
        normalize,
        layers.Dense(64),
        layers.Dense(1)
    ])
    norm_games_model.compile(
        loss=tensorflow.keras.losses.MeanSquaredError(),
        optimizer=tensorflow.keras.optimizers.Adam())

    norm_games_model.fit(games_all, games_predict, epochs=EPOCHS)

    norm_games_model.save('test')
    _run.add_artifact('test')

exint.run()
exint.add_artifact("test/saved_model.pb")
exint.add_source_file('./ium_z434686/train.py')