# import libraries
import numpy as np
import pandas as pd 
from matplotlib import pyplot as plt
%matplotlib inline
import seaborn as sns
import pandas as pd

# import in the training data set
train = pd.read_csv("train_features_2013-03-07.csv", low_memory=False)
# only drop rows that are all NA:
train = train.dropna(how='all')

# read the target variable - salary 
target = pd.read_csv("train_salaries_2013-03-07.csv")

# read the test data
test = pd.read_csv("test_features_2013-03-07.csv")

# take a look at the first 5 rows of train data
train.head()

# take a look at the first 5 rows of test data
test.head()

# take a look at the first 5 rows of target data
target.head()

# check the number of entries and data type for each variable
train.info()

# check the number of entries and data type for each variable
test.info()

# check the number of entries and data type for each variable
target.info()

## check duplicates in train data
print('Are all observations in train unique?')
print(len(train.jobId.unique()) == train.shape[0]) 

## check duplicates in test data
print('Are all observations in test unique?')
print(len(test.jobId.unique()) == test.shape[0]) 

# now check if there's any overlap between train and test in terms of job postings
print('The number of job IDs in the test dataset that did appear in the train dataset:')
print(len ( set(test['jobId']).intersection( train['jobId'] ) ) )

print('The number of job IDs in the train dataset that did appear in the test dataset:')
print(len ( set(train['jobId']).intersection( test['jobId'] ) ) )

# let's check if the jobID appear in the same order in train and in target
# if not, we'll need to sort the target to match the train dataset
sum(train.jobId == target.jobId) == train.shape[0]

sum(train.columns == test.columns) == train.shape[1]

# for the purpose of less typing, let's concat the train and test data
data = pd.concat([train, test], keys=['train','test'])
# get a quick description of numeric values in the data
data.describe()

# histogram of years of experience
sns.distplot(data.yearsExperience, kde= False, bins=25)

# histogram of miles from metropolis
sns.distplot(data.milesFromMetropolis, kde = False)

target.describe()

# count the number of 0s
print('The number of 0s in salary: '+ str( sum(target.salary == 0)))
print('The minimum non-zero salary: ' + str(min(target.salary[target.salary != 0])))

# remove zero-salary observations
zero_index = target[target.salary == 0].index
train.drop(zero_index, inplace=True)
target.drop(zero_index, inplace=True)

# update data
data = pd.concat([train, test], keys=['train','test'])
# describe the new data
target.describe()

data.select_dtypes(include=['object']).describe()

# plot the distribution of the target
sns.distplot(target.salary)

# get categorical variable names
cat_vars = train.select_dtypes(include=['object']).columns
for i in cat_vars:
    print('The variable ' + i + ' has the same levels in the train and the test dataset:')
    # check if the intersection of levels is the same set as in train data
    print(len ( set(train[i]).intersection( test[i] ) ) == len( set(train[i])) )

# first, for visualization purpose, create numeric representation for categorical variables 
train_cp = train.copy()
for i in train_cp.select_dtypes(include=['object']).columns.values:
    train_cp[i] = pd.factorize(train_cp[i])[0]
# add in the outcome var - salary 
train_cp['salary'] = target.salary
# remove job Id
train_cp.drop(['jobId'],axis=1, inplace=True)

# second, Create a pair grid instance
# since the plot with all the train data took very long to generate
# let's plot on a random subset of the data
grid = sns.PairGrid(data= train_cp.sample(frac= 0.1) )

# Map the plots to the locations
grid = grid.map_offdiag(sns.regplot, marker = '+', x_jitter=.1, y_jitter=0.1)
grid = grid.map_diag(sns.distplot, kde = False)

import h2o
h2o.init()

# convert pandas DataFrame to H2OFrame 
dat = train.merge(target, on = 'jobId')
df = h2o.H2OFrame(dat)
# split the data into train and test
splits = df.split_frame(ratios=[0.8], seed=42)
dtrain = splits[0]
dtest = splits[1]

y = 'salary'
x = list(df.columns)
x.remove(y)  #remove the response
x.remove('jobId') 
# The companyId feature doesn't seem to relate with salary, which comfirms our intuition. 
# Also in production, there's always the possibility of predicting on unseen companies before.
# Therefore, let's remove the companyId feature for now
x.remove('companyId') 

## grid search for best params with cross validation
from h2o.grid.grid_search import H2OGridSearch
from h2o.estimators.glm import H2OGeneralizedLinearEstimator
# grid search criteria, in the interest of time, let's set an upper limit to the runtime and max models
gs_criteria = {'strategy': 'RandomDiscrete', 
              'max_models': 100,
              'max_runtime_secs': 6000, 
              'stopping_metric': 'mae',
              'stopping_tolerance': 3,
              'seed': 1234,
              'stopping_rounds': 15}

# linear model parameter space
lm_params = {'alpha': [0.1, 0.5, 0.9], 
             'lambda': [1e-4, 1e-2, 1e0, 1e2, 1e4]}

lm_gs = H2OGridSearch(model = H2OGeneralizedLinearEstimator(nfolds = 5), 
                      hyper_params = lm_params,
                     search_criteria = gs_criteria)
lm_gs.train(x=x,y=y, training_frame=dtrain)
# display model IDs, hyperparameters, and MSE
# lm_gs.show()

# Get the grid results, sorted by validation mae
lm_gs_perf = lm_gs.get_grid(sort_by='mae', decreasing=True)
lm_gs_perf

# extract the top lm model, chosen by validation mae
lm_best = lm_gs.models[0]

# Now let's evaluate the model performance on a test set
# so we get an honest estimate of top model performance
lm_best_perf = lm_best.model_performance(dtest)

print('The best test MAE is %.2f' % lm_best_perf.mae())
print('The best test RMSE is %.2f' % lm_best_perf.rmse())
print('The best test R2 is %.2f' % lm_best_perf.r2())

# get coef
coef = lm_best.coef()
# plot feature importance
coef_df = pd.DataFrame({'feature':coef.keys(), 'coefficient':coef.values()})

coef_df.reindex(coef_df.coefficient.abs().sort_values(ascending = False).index)

# now let's use the random forest algorithm to see if we can improve the performance
# import 
from h2o.estimators.random_forest import H2ORandomForestEstimator
# rf parameters to be tuned
rf_params = {"ntrees" : [50,100], 
              "max_depth" : [5,10], 
              "sample_rate" : [0.3, 0.7], 
              "mtries":[-1, 3, 7],
              "min_rows" : [3,5] }
# parameters for rf that don't need to be tuned
rf_p = {"seed": 42 , 
         "score_tree_interval": 3,
         "nfolds": 5}
# set up the grid
rf_gs = H2OGridSearch(model = H2ORandomForestEstimator(**rf_p), 
                       hyper_params = rf_params,
                       search_criteria = gs_criteria)
# train the grid
rf_gs.train(x=x,y=y, training_frame=df)

# display model IDs, hyperparameters, and MSE
# rf_gs.show()

# Get the grid results, sorted by validation mae
rf_gs_perf = rf_gs.get_grid(sort_by='mae', decreasing=True)
rf_gs_perf

# extract the top lm model, chosen by validation mae
rf_best = rf_gs.models[0]

# Now let's evaluate the model performance on a test set
# so we get an honest estimate of top model performance
rf_best_perf = rf_best.model_performance(dtest)

print('The best test MAE is %.2f' % rf_best_perf.mae())
print('The best test RMSE is %.2f' % rf_best_perf.rmse())
print('The best test R2 is %.2f' % rf_best_perf.r2())

# plot feature importance
rf_best.varimp_plot()

# write output
# convert test data to h2o frame
test_h2o = h2o.H2OFrame(test)
# determine the best model based on mae and use it for prediction
if rf_best_perf.mae() < lm_best_perf.mae():
    pred = rf_best.predict(test_h2o)
else:
    pred = lm_best.predict(test_h2o)

# examine the basic stats of the salary predictions
pred.describe()

# concat jobId and salary prediction
out = test_h2o['jobId'].concat(pred)
# change names
out.set_names(list(target.columns))
# write the test output to a csv file
h2o.export_file(out, path = "test_salaries.csv")