#!/usr/bin/env python3
'''
This code takes the previously created sample sets and creates a model to
'''
import os
import tensorflow as tf
import pandas as pd


X_train = pd.read_csv('support.csv')
y_train = pd.read_csv('support_labels.csv')

NUM_TRAINING_RUNS = 1
NUM_TRAINING_EPOCHS = 3


print(f"{X_train.shape=}")
print(f"{y_train.shape=}")

model = tf.keras.Sequential([
    tf.keras.layers.Dense(1024, activation='relu'),
    tf.keras.layers.Dense(2048, activation='relu'),
    tf.keras.layers.Dense(2048, activation='relu'),
    tf.keras.layers.Dense(1, activation='sigmoid')
])

model.compile(
    loss=tf.keras.losses.binary_crossentropy,
    optimizer=tf.keras.optimizers.Adam(learning_rate=0.003),
    metrics=[
        tf.keras.metrics.BinaryAccuracy(name='accuracy'),
        tf.keras.metrics.Precision(name='precision'),
        tf.keras.metrics.Recall(name='recall')
    ]
    )

model.fit(X_train, y_train, epochs=NUM_TRAINING_EPOCHS)

model.summary()

save_path = "support_model"
if not os.path.exists(save_path):
    os.mkdir(save_path)

tf.saved_model.save(model,save_path)

print(f"Loading saved model from {save_path} and creating converter")
converter = tf.lite.TFLiteConverter.from_saved_model(save_path)

#converter = tf.lite.TFLiteConverter.from_keras_model(model)


print("Calling converter.convert()")
tflite_model = converter.convert()
print("Writing output")
with open(os.path.join("save_path", 'model.tflite'), 'wb') as f:
    f.write(tflite_model)