import sys
from tensorflow.keras import layers, Sequential
# from keras.layers import Flatten,Dense,Dropout, GlobalAveragePooling2D
from tensorflow.keras.optimizers import Adam
from tensorflow import convert_to_tensor
import numpy as np
import pandas as pd
from sklearn.metrics import mean_squared_error
from sacred.observers import FileStorageObserver, MongoObserver
from sacred import Experiment
from sacred.observers import MongoObserver
from datetime import datetime


ex = Experiment("434684", 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())

	movies_train = pd.read_csv('movies_train.csv')

	x_train = movies_train.copy()
	y_train = x_train.pop('rottentomatoes_audience_score')
	x_train.pop('Unnamed: 0')

	learning_rate = float(sys.argv[1])
	model = Sequential()
	model.add(layers.Input(shape=(22,)))
	model.add(layers.Dense(64))
	model.add(layers.Dense(64))
	model.add(layers.Dense(32))
	model.add(layers.Dense(1))

	model.compile(loss='mean_absolute_error', optimizer=Adam(learning_rate))
	
	_run.info["train model"] = str(datetime.now())
	
	history = model.fit(
		x = convert_to_tensor(x_train, np.float32),
		y = y_train,
		verbose=0, epochs=99)
		
	loss = history.history['loss'][-1]
	_run.info["Loss"] = str(loss)
	model.save('model_movies.h5')
		

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


r = ex.run()
ex.add_artifact("model_movies.h5")