from tensorflow.keras.models import Sequential, load_model
from tensorflow.keras.layers import Dense
from sklearn.metrics import accuracy_score, classification_report
import pandas as pd
from urllib.parse import urlparse
from sklearn.model_selection import train_test_split
import numpy as np
import sys
from sklearn.preprocessing import StandardScaler, LabelEncoder
from tensorflow.keras.optimizers import Adam
from datetime import datetime
import os
import sys
import mlflow

mlflow.set_tracking_uri("http://172.17.0.1:5000")
mlflow.set_experiment("s434788")

with mlflow.start_run():

    batch_param = int(sys.argv[1]) if len(sys.argv) > 1 else 16
    epoch_param = int(sys.argv[2]) if len(sys.argv) > 2 else 15
   
    
    mlflow.log_param("batch_size", batch_param)
    mlflow.log_param("epochs", epoch_param)

    wine=pd.read_csv('winequality-red.csv')

    y = wine['quality']
    x = wine.drop('quality', axis=1)

    citricacid = x['fixed acidity'] * x['citric acid']
    citric_acidity = pd.DataFrame(citricacid, columns=['citric_accidity'])

    density_acidity = x['fixed acidity'] * x['density']
    density_acidity = pd.DataFrame(density_acidity, columns=['density_acidity'])


    x = wine.join(citric_acidity).join(density_acidity)

    bins = (2, 5, 8)
    gnames = ['bad', 'nice']
    y = pd.cut(y, bins = bins, labels = gnames)

    enc = LabelEncoder()
    yenc = enc.fit_transform(y)

    scale = StandardScaler()
    scaled_x = scale.fit_transform(x)
    
    x_train, x_test, y_train, y_test = train_test_split(scaled_x,yenc , test_size=0.2,train_size=0.8, random_state=21)

    NeuralModel = Sequential([
                            Dense(128, activation='relu', input_shape=(14,)),
                            Dense(32, activation='relu'),
                            Dense(64, activation='relu'),
                            Dense(64, activation='relu'),
                            Dense(64, activation='relu'),
                            Dense(1, activation='sigmoid')
    ])

    rms = Adam(lr=0.0003)

    NeuralModel.compile(optimizer=rms, loss='binary_crossentropy', metrics=['accuracy'])

    NeuralModel.fit(x_train, y_train, batch_size= batch_param, epochs = epoch_param) #verbose = 1

    #signature = mlflow.models.signature.infer_signature(x_train.values, NeuralModel.predict(x_train.values))
    #input_example = x_test.values[10]

    y_pred = NeuralModel.predict(x_test)

    y_pred = np.around(y_pred, decimals=0)

    results = accuracy_score(y_test,y_pred)

    print(f"accuracy: {results}")

    mlflow.log_metric("Accuracy", results)
    
    tracking_url_type_store = urlparse(mlflow.get_tracking_uri()).scheme

    if tracking_url_type_store != "file":
        mlflow.keras.save_model(NeuralModel, "my_model", signature= mlflow.models.signature.infer_signature(x_train, y_train), input_example= x_train[0])
        mlflow.keras.log_model(NeuralModel, "model",registered_model_name="s434788", signature= mlflow.models.signature.infer_signature(x_train, y_train), input_example= x_train[0])
    else:
        mlflow.keras.log_model(NeuralModel, "model", signature= mlflow.models.signature.infer_signature(x_train, y_train))