# The data set used in this example is from http://archive.ics.uci.edu/ml/datasets/Wine+Quality
# P. Cortez, A. Cerdeira, F. Almeida, T. Matos and J. Reis.
# Modeling wine preferences by data mining from physicochemical properties. In Decision Support Systems, Elsevier, 47(4):547-553, 2009.

# Modified from `train.py` to be run interactively
# Zhenhao Ge, 2019-05-16

import os, sys
import warnings
import __main__ as main

import pandas as pd
import numpy as np
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
from sklearn.model_selection import train_test_split
from sklearn.linear_model import ElasticNet

import mlflow
import mlflow.sklearn

def eval_metrics(actual, pred):
    rmse = np.sqrt(mean_squared_error(actual, pred))
    mae = mean_absolute_error(actual, pred)
    r2 = r2_score(actual, pred)
    return rmse, mae, r2

if __name__ == '__main__':
    warnings.filterwarnings('ignore')
    np.random.seed(40)

    # Read the wine-quality csv file (make sure you're running this from the root of MLflow!)
    # '__file__' shows the relative path of current running script
    # e.g. sklearn_elasticnet_wine/train.py, if the working dir is mlflow/examples
    if hasattr(main, '__file__'):
      print('running mode ...')
      wine_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'wine-quality.csv')
    else:
      print('interactive mode ...')
      work_dir = r'/home/zge/Work/Projects/mlflow/examples/sklearn_elastic_wine'
      os.chdir(work_dir)
      wine_path = os.path.join(work_dir, 'sklearn_elasticnet_wine', 'wine-quality.csv')
    #print('wine path: {}'.format(wine_path))

    data = pd.read_csv(wine_path)

    # Split the data into training and test sets. (0.75, 0.25) split.
    train, test = train_test_split(data, train_size=0.75, random_state=0)

    # The predicted column is "quality" which is a scalar from [3, 9]
    train_x = train.drop(["quality"], axis=1)
    test_x = test.drop(["quality"], axis=1)
    train_y = train[["quality"]]
    test_y = test[["quality"]]

    alpha = float(sys.argv[1]) if len(sys.argv) > 1 else 0.5
    l1_ratio = float(sys.argv[2]) if len(sys.argv) > 2 else 0.5
    print('alpha: {}, l1_ratio: {}'.format(alpha, l1_ratio))

    with mlflow.start_run():
        # alpha: rugulization factor (the larger, the more regularization)
        # l1_ratio (percentage of l1-norm regularization)
        lr = ElasticNet(alpha=alpha, l1_ratio=l1_ratio, random_state=42)
        lr.fit(train_x, train_y)

        predicted_qualities = lr.predict(test_x)

        (rmse, mae, r2) = eval_metrics(test_y, predicted_qualities)

        print("Elasticnet model (alpha=%f, l1_ratio=%f):" % (alpha, l1_ratio))
        print("  RMSE: %s" % rmse)
        print("  MAE: %s" % mae)
        print("  R2: %s" % r2)

        mlflow.log_param("alpha", alpha)
        mlflow.log_param("l1_ratio", l1_ratio)
        mlflow.log_metric("rmse", rmse)
        mlflow.log_metric("r2", r2)
        mlflow.log_metric("mae", mae)

        mlflow.sklearn.log_model(lr, "model")
