import logging
import os
import os.path
import shutil
import tempfile

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import requests
import scipy.stats as stats
import seaborn as sns
from sklearn.model_selection import train_test_split

from ludwig.api import kfold_cross_validate, LudwigModel

"""
Explanation: K-fold cross validation - Regression Model
Based on the Ludwig regression example  
Data set
This example demonstrates teh following:

Download a data set and create a pandas dataframe
Create a training and hold-out test data sets
Create a Ludwig config data structure from the pandas dataframe
Run a 5-fold cross validation analysis with the training data
Use Ludwig APIs to train and assess model performance on hold-out test data set
End of explanation
"""


DATA_SET_URL = 'http://archive.ics.uci.edu/ml/machine-learning-databases/auto-mpg/auto-mpg.data'
DATA_SET = 'auto_mpg.data'
RESULTS_DIR = 'results'

"""
Explanation: Contstants
End of explanation
"""


if os.path.isfile(DATA_SET):
    os.remove(DATA_SET)
    
shutil.rmtree(RESULTS_DIR, ignore_errors=True)

"""
Explanation: Clean out previous results
End of explanation
"""


r = requests.get(DATA_SET_URL)
if r.status_code == 200:
    with open(DATA_SET,'w') as f:
        f.write(r.content.decode("utf-8"))

"""
Explanation: Retrieve data from UCI Machine Learning Repository
Download required data
End of explanation
"""


raw_df = pd.read_csv(DATA_SET,
                     header=None,
                      na_values = "?", comment='\t',
                      sep=" ", skipinitialspace=True)


raw_df.columns = ['MPG','Cylinders','Displacement','Horsepower','Weight',
                'Acceleration', 'ModelYear', 'Origin']
raw_df.shape

raw_df.head()

"""
Explanation: Create Pandas DataFrame from downloaded data
End of explanation
"""


train_df, test_df = train_test_split(raw_df, train_size=0.8, random_state=17)
print(train_df.shape)
print(test_df.shape)

"""
Explanation: Create train/test split
End of explanation
"""


num_features = ['Cylinders', 'Displacement', 'Horsepower', 'Weight', 'Acceleration', 'ModelYear']
cat_features = ['Origin']

"""
Explanation: Setup Ludwig config
End of explanation
"""


input_features = []
# setup input features for numerical variables
for p in num_features:
    a_feature = {'name': p, 'type': 'numerical', 
                'preprocessing': {'missing_value_strategy': 'fill_with_mean', 'normalization': 'zscore'}}
    input_features.append(a_feature)

# setkup input features for categorical variables
for p in cat_features:
    a_feature = {'name': p, 'type': 'category'}

"""
Explanation: Create Ludwig input_features
End of explanation
"""


output_features =[
    {
        'name': 'MPG',
        'type': 'numerical',
        'num_fc_layers': 2,
        'fc_size': 64
    }
]

config = {
    'input_features' : input_features,
    'output_features': output_features,
    'training' :{
        'epochs': 100,
        'batch_size': 32
    }
}
config

"""
Explanation: Create Ludwig output features
End of explanation
"""


%%time
with tempfile.TemporaryDirectory() as tmpdir:
    data_csv_fp = os.path.join(tmpdir,'train.csv')
    train_df.to_csv(data_csv_fp, index=False)

    (
        kfold_cv_stats, 
        kfold_split_indices 
    ) = kfold_cross_validate(
        num_folds=5,
        config=config,
        dataset=data_csv_fp,
        output_directory=tmpdir,
        logging_level=logging.ERROR
    )


kfold_cv_stats['overall']['MPG']

"""
Explanation: Perform K-fold Cross Validation analysis
End of explanation
"""


model = LudwigModel(
    config=config,
    logging_level=logging.ERROR
)

%%time
training_stats = model.train(
    training_set=train_df,
    output_directory=RESULTS_DIR,
)

test_stats, mpg_hat_df, _ = model.evaluate(dataset=test_df, collect_predictions=True, collect_overall_stats=True)

test_stats

a = plt.axes(aspect='equal')
sns.scatterplot(test_df['MPG'].values, mpg_hat_df['MPG_predictions'].values,
               s=50)
plt.xlabel('True Values [MPG]')
plt.ylabel('Predictions [MPG]')
lims = [0, 50]
plt.xlim(lims)
plt.ylim(lims)
_ = plt.plot(lims, lims)


"""
Explanation: Train model and assess model performance
End of explanation
"""


test_stats['MPG']

"""
Explanation: Compare K-fold Cross Validation metrics against hold-out test metrics
Hold-out Test Metrics
End of explanation
"""


kfold_cv_stats['overall']['MPG']

"""
Explanation: K-fold Cross Validation Metrics
End of explanation
"""


