import numpy as np
import pandas as pd
import tensorflow as tf
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
import sys

def main():
    no_of_epochs = 10
    if len(sys.argv) == 2:
        no_of_epochs = int(sys.argv[1])

    scaler = StandardScaler()
    feature_names = ["BMI", "SleepTime", "Sex", "Diabetic", "PhysicalActivity", "Smoking", "AlcoholDrinking"]

    dataset = pd.read_csv('heart_2020_cleaned.csv')
    dataset = dataset.dropna()

    dataset["Diabetic"] = dataset["Diabetic"].apply(lambda x: True if "Yes" in x else False)
    dataset["HeartDisease"] = dataset["HeartDisease"].apply(lambda x: True if x == "Yes" else False)
    dataset["PhysicalActivity"] = dataset["PhysicalActivity"].apply(lambda x: True if x == "Yes" else False)
    dataset["Smoking"] = dataset["Smoking"].apply(lambda x: True if x == "Yes" else False)
    dataset["AlcoholDrinking"] = dataset["AlcoholDrinking"].apply(lambda x: True if x == "Yes" else False)
    dataset["Sex"] = dataset["Sex"].apply(lambda x: 1 if x == "Female" else 0)

    dataset_train, dataset_test = train_test_split(dataset, test_size=.1, train_size=.9, random_state=1)

    print(dataset_test.shape)

    model = tf.keras.Sequential([
        tf.keras.layers.Dense(16, activation='relu'),
        tf.keras.layers.Dense(8, activation='relu'),
        tf.keras.layers.Dense(4, activation='relu'),
        tf.keras.layers.Dense(1, activation='sigmoid')
    ])

    model.compile(
        loss=tf.keras.losses.binary_crossentropy,
        optimizer=tf.keras.optimizers.Adam(lr=0.01),
        metrics=["accuracy", tf.keras.metrics.Recall(name='recall')]
    )

    train_X = dataset_train[feature_names].astype(np.float32)
    train_Y = dataset_train["HeartDisease"].astype(np.float32)
    test_X = dataset_test[feature_names].astype(np.float32)
    test_Y = dataset_test["HeartDisease"].astype(np.float32)

    train_X = scaler.fit_transform(train_X)
    # train_Y = scaler.fit_transform(train_Y)
    test_X = scaler.fit_transform(test_X)
    # test_Y = scaler.fit_transform(test_Y)


    print(train_Y.value_counts())
    train_X = tf.convert_to_tensor(train_X)
    train_Y = tf.convert_to_tensor(train_Y)

    test_X = tf.convert_to_tensor(test_X)
    test_Y = tf.convert_to_tensor(test_Y)

    model.fit(train_X, train_Y, epochs=no_of_epochs)
    model.save("trained_model")

    test_loss, test_acc, test_rec = model.evaluate(test_X, test_Y, verbose=1)
    print("Accuracy:", test_acc)
    print("Loss:", test_loss)
    print("Recall:", test_rec)


if __name__ == '__main__':
    main()