# Import the necessary packages
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
%matplotlib inline
plt.style.use('ggplot')

# Where to save the figures
PROJECT_ROOT_DIR = "../.."
datapath = PROJECT_ROOT_DIR + "/data/lifesat/"

"""
Explanation: Intro to SML
This is based on the notebook file 01 in Aurélien Geron's github page
End of explanation
"""


# Download CSV from http://stats.oecd.org/index.aspx?DataSetCode=BLI
oecd_bli = pd.read_csv(datapath+"oecd_bli_2015.csv", thousands=',')
oecd_bli = oecd_bli[oecd_bli["INEQUALITY"]=="TOT"]
oecd_bli.head()

oecd_bli = oecd_bli.pivot(index="Country", columns="Indicator", values="Value")
oecd_bli.head()

oecd_bli.columns

oecd_bli["Life satisfaction"].head()

"""
Explanation: Load and prepare Life satisfaction data
End of explanation
"""


# Download data from http://goo.gl/j1MSKe (=> imf.org)
gdp_per_capita = pd.read_csv(datapath+"gdp_per_capita.csv", thousands=',', delimiter='\t',
                             encoding='latin1', na_values="n/a")
gdp_per_capita.rename(columns={"2015": "GDP per capita"}, inplace=True)
gdp_per_capita.set_index("Country", inplace=True)
gdp_per_capita.head()

full_country_stats = pd.merge(left=oecd_bli, right=gdp_per_capita, left_index=True, right_index=True)
full_country_stats.sort_values(by="GDP per capita", inplace="True")
full_country_stats.head()

full_country_stats[["GDP per capita", 'Life satisfaction']].loc["United States"]

remove_indices = [0, 1, 6, 8, 33, 34, 35]
keep_indices = list(set(range(36)) - set(remove_indices))

sample_data = full_country_stats[["GDP per capita", 'Life satisfaction']].iloc[keep_indices]
missing_data = full_country_stats[["GDP per capita", 'Life satisfaction']].iloc[remove_indices]

_ = sample_data.plot("GDP per capita",'Life satisfaction',kind='scatter')

"""
Explanation: Load and prepare GDP per capita data
End of explanation
"""


x = np.array(sample_data[['GDP per capita']])
n,_ = x.shape
X = np.hstack((np.ones((n,1)),x))
y = np.array(sample_data[['Life satisfaction']])

print(X[0:5,:])

"""
Explanation: Linear regression to predict Life satisfaction using GDP per capita
Let's make a design matrix and response vector.
End of explanation
"""


from sklearn import linear_model
lin1 = linear_model.LinearRegression(fit_intercept=False)
lin1.fit(X,y)
print(lin1.coef_)

yhat = lin1.predict(X)
# Compute the training error
RSS = np.sum((y - yhat)**2.)
training_error = RSS / n
print(training_error)

_ = plt.scatter(X[:,1],y)
_ = plt.plot(X[:,1],yhat,'r')

"""
Explanation: Ordinary least squares is defined by the normal equation
$$X^\top X \hat \beta = X^\top y$$
$$\hat y = X \hat \beta$$
Let's fit the linear model and predict the points.
End of explanation
"""


X_quad = np.hstack((np.ones((n,1)),x,x**2))

"""
Explanation: The regression line seems to demonstrate a good fit to a linear model.  We can use it to predict a new point, to assess the impact of a tax policy for example.
Quadratic Regression
End of explanation
"""


from sklearn import linear_model
lin1 = linear_model.LinearRegression(fit_intercept=False)
lin1.fit(X_quad,y)
yhat = lin1.predict(X_quad)
# Compute the training error
RSS = np.sum((y - yhat)**2.)
training_error = RSS / n
print(training_error)

_ = plt.scatter(X[:,1],y)
_ = plt.plot(X[:,1],yhat,'r')

"""
Explanation: In the above, we add a quadratic term to the design matrix, to see if it improves the fit.
End of explanation
"""


from sklearn.model_selection import LeaveOneOut
loo = LeaveOneOut()
loo.get_n_splits(X)

loo_losses = []
for train_index, test_index in loo.split(X):
    X_train, X_test = X[train_index], X[test_index]
    y_train, y_test = y[train_index], y[test_index]
    lin1.fit(X_train,y_train)
    y_hat = lin1.predict(X_test)
    loss = np.sum((y_hat - y_test)**2)
    loo_losses.append(loss)
    print("TEST: " + sample_data.index.values[test_index][0] + " " + str(loss))

plt.hist(loo_losses)
plt.vlines(training_error,0,12,'b',label="Training Risk")
plt.vlines(np.mean(loo_losses),0,12,'k',label="LOO Ave Risk")
_ = plt.legend()

"""
Explanation: The fits are essentially the same between linear and quadratic fit, and the training risk is nearly equal.
Leave-one-out prediction simulation
In this section, I show that the training risk is an optimistic measure of the true risk.
End of explanation
"""


def loo_risk(X,y):
    """
    Construct the leave-one-out square error risk for the linear model
    
    Input: design matrix, X, response vector, y
    Output: scalar LOO risk
    """
    loo = LeaveOneOut()
    loo.get_n_splits(X)
    loo_losses = []
    for train_index, test_index in loo.split(X):
        X_train, X_test = X[train_index], X[test_index]
        y_train, y_test = y[train_index], y[test_index]
        lin1.fit(X_train,y_train)
        y_hat = lin1.predict(X_test)
        loss = np.sum((y_hat - y_test)**2)
        loo_losses.append(loss)
    return np.mean(loo_losses)

help(loo_risk)

print("Linear LOO risk: {}".format(loo_risk(X,y)))
print("Quadratic LOO risk: {}".format(loo_risk(X_quad,y)))

"""
Explanation: This histogram of the leave-one-out errors demonstrates that the training risk can be optimistic, because the Leave-one-out risk (the average of the LOO errors) is larger than the training risk.
End of explanation
"""


