import pandas as pd
import numpy as np
import tensorflow as tf
import os.path

import mlflow
import sys

from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.layers.experimental import preprocessing

arguments = sys.argv[1:]

verbose = int(arguments[0])
epochs = int(arguments[1])

 # Wczytanie danych
train_data = pd.read_csv("./MoviesOnStreamingPlatforms_updated.train")
test_data = pd.read_csv("./MoviesOnStreamingPlatforms_updated.test")

# Stworzenie modelu 
columns_to_use = ['Year', 'Runtime', 'Netflix']
train_X = tf.convert_to_tensor(train_data[columns_to_use])
train_Y = tf.convert_to_tensor(train_data[["IMDb"]])

test_X = tf.convert_to_tensor(test_data[columns_to_use])
test_Y = tf.convert_to_tensor(test_data[["IMDb"]])

normalizer = preprocessing.Normalization(input_shape=[3,])
normalizer.adapt(train_X)

model = keras.Sequential([
    keras.Input(shape=(len(columns_to_use),)),
    normalizer,
    layers.Dense(30, activation='relu'),
    layers.Dense(10, activation='relu'),
    layers.Dense(25, activation='relu'),
    layers.Dense(1)
])

model.compile(loss='mean_absolute_error',
                optimizer=tf.keras.optimizers.Adam(0.001),
                metrics=[tf.keras.metrics.RootMeanSquaredError()])

model.fit(train_X, train_Y, verbose=verbose, epochs=epochs)

signature = mlflow.models.signature.infer_signature(train_X.numpy(), model.predict(train_X.numpy()))
input_data = test_X

with mlflow.start_run():
    mlflow.keras.save_model(model, "movies_on_streaming_platforms_model", input_example=input_data.numpy(), signature=signature)