# Dependencies
%config InlineBackend.figure_format = 'retina'
%matplotlib inline

import numpy as np
import matplotlib.pyplot as plt

"""
Explanation: Homework 2
In this homework you are going to fit and compare linear models on how well they predict new data. You will need to install (at least) python, jupyter, matplotlib, and numpy to do this assignment. Installing anaconda is a quick way to get all of those things.
There are 2 problems worth a total of 30 points. The description for each problem will tell you how many points each part is worth.
You should not need to edit the boilerplate code in this notebook, but wherever you see ## YOUR CODE HERE ##, you should replace that with your own code (obviously).
To turn in your homework, email your finished .ipynb file to huth@cs.utexas.edu. This homework is due on 11/14.
End of explanation
"""


# Load data!
p1a_file = np.load('homework_2_p1a_data.npz')
x = p1a_file['x']
y = p1a_file['y']
x_test = p1a_file['x_test']
y_test = p1a_file['y_test']

n_training, n_features = x.shape

"""
Explanation: Problem 1 - Cross-validation (10 pts)
Cross-validation is the method that we can use to choose hyperparameters, such as the ridge penalty $\lambda$ in ridge regression. It is an empirical technique that is incredibly effective, even if the data does not exactly fit the assumptions of your model (e.g. if the errors are not gaussian, etc.). You can read more about cross-validation [here](https://en.wikipedia.org/wiki/Cross-validation_(statistics%29).
As a basic scheme, imagine that you have data that is split into two parts: training and test. Let's suppose you try a bunch of different ridge parameters by fitting models using the training data. If you check how well each model predicts the training data (that you used to generate the weights), then you'll get a bad answer: most likely the smallest ridge parameter will work best (you saw this in the last homework). If you check how well each model predicts the test data, then you're cheating: you're touching the test data before you're done model fitting. Don't cheat by fitting hyperparameters using the test data!! So obviously we need another way to check the ridge parameters.
Now suppose that you further split the training dataset into two parts: the "training" set (yes it's the same name, get over it), and a "validation" set. Now you can train your models using the training data, check how well they work using the validation dataset, and use those numbers to pick your hyperparameter. Then, finally, you can use the whole training dataset and the best hyperparameter to fit a model that you'll use to try to predict the test data (which you haven't touched until now). No cheating! This is cross-validation.
In this problem, you are going to use three different cross-validation methods to select $\lambda$ for ridge regression.
BIG HINT: The np.delete function is your best friend for this problem.
End of explanation
"""


lambdas = np.logspace(-3, 5, 10)
val_ses = np.zeros((n_training, len(lambdas)))

def ridge(x, y, lam):
    """This function does ridge regression with the stimuli x and responses y with
    ridge parameter lam (short for lambda). It returns the weights.
    This is definitely not the most efficient way to do this, but it's fine for now.
    """
    n_features = x.shape[1]
    beta_ridge = np.linalg.inv(x.T.dot(x) + lam * np.eye(n_features)).dot(x.T).dot(y)
    return beta_ridge

for t in range(n_training):
    # split the training dataset into two parts: one with only point t,
    # and one with all the other datapoints
    x_trn = ## YOUR CODE HERE ##
    y_trn = ## YOUR CODE HERE ##
    
    x_val = ## YOUR CODE HERE ##
    y_val = ## YOUR CODE HERE ##
    
    for ii in range(len(lambdas)):
        # fit model using x_trn & predict y_hal
        y_val_hat = ## YOUR CODE HERE ##
        
        # store squared error in val_mses
        val_ses[t,ii] = ## YOUR CODE HERE ##

# Plot the mean squared error with each lambda, averaged across validation datapoints

## YOUR CODE HERE ##

# Choose the best lambda (i.e. the one with the least validation error), print it
best_lambda = ## YOUR CODE HERE ##
print("Best lambda:", best_lambda)

# Fit a model using the whole training set and the best lambda
beta_hat = ## YOUR CODE HERE ##

# Use that model to predict y_test
y_test_hat = ## YOUR CODE HERE ##

# Compute the MSE on the test dataset, print it
test_mse = ## YOUR CODE HERE ##

print("Test MSE:", test_mse)

"""
Explanation: (a) Leave-one-out (LOO) cross-validation (3 pts)
This is perhaps the simplest cross-validation method. Select one datapoint from the training set to be the validation set, then train on the other $n-1$ datapoints. Test on the remaining datapoint. Repeat this process for each datapoint, and then average the results.
This works very well when the datapoints are independent, which is going to be the case in this first dataset.
End of explanation
"""


k = 6 # let's do 6 folds
n_per_fold = n_training / k # number of datapoints per fold

lambdas = np.logspace(-3, 5, 10)
val_mses = np.zeros((k, len(lambdas)))


for fold in range(k):
    # split the training dataset into two parts: one with only the points in fold "fold"
    # and one with all the other datapoints
    
    ## YOUR CODE HERE ##
    
    x_trn = ## YOUR CODE HERE ##
    y_trn = ## YOUR CODE HERE ##
    
    x_val = ## YOUR CODE HERE ##
    y_val = ## YOUR CODE HERE ##
    
    for ii in range(len(lambdas)):
        # fit model using x_trn & predict y_val
        y_val_hat = ## YOUR CODE HERE ##
        
        # store squared error in val_mses
        val_mses[fold,ii] = ## YOUR CODE HERE ##

# Plot the MSE for each lambda, averaged across the folds
## YOUR CODE HERE ##

# Choose the best lambda, print it
best_lambda = ## YOUR CODE HERE ##
print("Best lambda:", best_lambda)

# Fit a model using the whole training set and the best lambda
beta_hat = ## YOUR CODE HERE ##

# Use that model to predict y_test
y_test_hat = ## YOUR CODE HERE ##

# Compute MSE on the test dataset, print it
test_mse = ## YOUR CODE HERE ##

print("Test MSE:", test_mse)

"""
Explanation: (b) $k$-fold cross-validation (3 pts)
You may have noticed that LOO CV is quite slow. That's because it needs to fit one model per datapoint. If you have a lot of datapoint, that's just not gonna work! Further, as mentioned above, LOO only works well when your datapoints are independent. Suppose you're recording fMRI data. You know that the underlying BOLD signal is very low frequency--meaning that the individual datapoints are anything but independent.
Another option is $k$-fold cross-validation. In this scheme, you split the training dataset into $k$ parts. Then you use $k-1$ of the parts to train the model, and use the $k$'th part as your validation dataset. Repeat this, holding out each part in turn. This means you only need to fit $k$ sets of models, instead of one for each datapoint.
If your entire training dataset has $n$ datapoints, each "fold" should contain $\frac{n}{k}$ datapoints, and each datapoint should only be in one fold. Assuming that $n/k$ is an integer, a nice way to do this is to break up the training dataset into $k$ contiguous parts.
End of explanation
"""


n_mc_iters = 50 # let's do 50 Monte Carlo iterations
n_per_mc_iter = 50 # on each MC iteration, hold out 50 datapoints to be the validation set

lambdas = np.logspace(-3, 5, 10)
val_mses = np.zeros((n_training, len(lambdas)))


for it in range(n_mc_iters):
    # split the training dataset into two parts: one with a random selection of n_per_mc_iter points
    # and one with all the other datapoints
    
    ## YOUR CODE HERE ##
    
    x_trn = ## YOUR CODE HERE ##
    y_trn = ## YOUR CODE HERE ##
    
    x_val = ## YOUR CODE HERE ##
    y_val = ## YOUR CODE HERE ##
    
    for ii in range(len(lambdas)):
        # fit model using x_trn & predict y_val
        # predict y_val
        y_val_hat = ## YOUR CODE HERE ##
        
        # store squared error in val_mses
        val_mses[it,ii] = ## YOUR CODE HERE ##

# Plot the MSE for each lambda, averaged across the MC iterations

## YOUR CODE HERE ##

# Choose the best lambda, print it
best_lambda = ## YOUR CODE HERE ##
print("Best lambda:", best_lambda)

# Fit a model using the whole training set and the best lambda
beta_hat = ## YOUR CODE HERE ##

# Use that model to predict y_test
y_test_hat = ## YOUR CODE HERE ##

# Compute the MSE, print it
test_mse = ## YOUR CODE HERE ##

print("Test MSE:", test_mse)

"""
Explanation: (c) Monte Carlo cross-validation (4 pts)
One issue with $k$-fold CV is that the size of the validation set depends on the number of folds. If you want really stable estimates for your hyperparameter, you want to have a pretty large validation set, but also do a lot of folds. You can accomplish this by, on each iteration, randomly assigning some fraction of the training set to be the validation set.
End of explanation
"""


from homework_2_utils import make_data

num_training = 500 # total number of datapoints in training set
num_test = 100 # total number of datapoints in test set
num_features = [12, 50, 100] # number of features in each feature space
num_responses = 35 # number of responses (voxels or neurons)

# This is just a bunch of constans. don't worry about what they mean for now
combs = [[0, 3, 4, 6],
         [1, 3, 5, 6],
         [2, 4, 5, 6]]
true_variances = np.array([300, 0, 1500, 250, 250, 4000, 500]).astype(float)
total_variance = 0.3
true_variances = true_variances / true_variances.sum() * total_variance
noise_variance = 1 - true_variances.sum()
P_parts = [3] * 7
Pnoise_models = [P - np.array(P_parts)[c].sum() for P,c in zip(num_features, combs)]

# Generate the data!
[x1_total, x2_total, x3_total], Y_total = make_data(num_training, num_test, P_parts, 
                                                    num_responses, true_variances, 
                                                    noise_variance, combs, Pnoise_models, 
                                                    num_features)

x1 = x1_total.T[:num_training]
x2 = x2_total.T[:num_training]
x3 = x3_total.T[:num_training]
Y = Y_total[:num_training]

x1_test = x1_total.T[num_training:]
x2_test = x2_total.T[num_training:]
x3_test = x3_total.T[num_training:]
Y_test = Y_total[num_training:]

print('x1 shape:', x1.shape)
print('x2 shape:', x2.shape)
print('x3 shape:', x3.shape)
print('Y shape:', Y.shape)

print('x1_test shape:', x1_test.shape)
print('x2_test shape:', x2_test.shape)
print('x3_test shape:', x3_test.shape)
print('Y_test shape:', Y_test.shape)


"""
Explanation: Problem 2 - Multiple feature spaces (20 pts)
Suppose you've done an experiment and measured responses to a bunch of stimuli. You've got three different hypotheses about how the stimuli might be represented in the responses. You instantiate these hypotheses as three different linearizing transforms, giving you three different sets of features that you can extract: $X_1$, $X_2$, and $X_3$ (in variables called x1, x2, and x3).
Feature space $X_1$ has 12 features, $X_2$ has 50 features, and $X_3$ has 100 features.
Note that you've recorded $m=35$ different responses (i.e. $Y$ is an $n\times m$ matrix). Think of these as $m$ different neurons or $m$ different voxels.
End of explanation
"""


lambdas = np.logspace(1, 7, 15) # use these lambdas

## YOUR CODE HERE ##

# Plot mean validation set MSE as a function of lambda for each feature space

## YOUR CODE HERE ##

# Find best lambda for each feature space
best_lambda_1 = ## YOUR CODE HERE ##
best_lambda_2 = ## YOUR CODE HERE ##
best_lambda_3 = ## YOUR CODE HERE ##
print best_lambda_1, best_lambda_2, best_lambda_3

# Fit models with best lambdas and entire training set
beta_1 = ## YOUR CODE HERE ##
beta_2 = ## YOUR CODE HERE ##
beta_3 = ## YOUR CODE HERE ##

# Predict Y_test with each, compute MSE
mse = lambda a, b: ((a - b)**2).mean()

test_mse_1 = ## YOUR CODE HERE ##
test_mse_2 = ## YOUR CODE HERE ##
test_mse_3 = ## YOUR CODE HERE ##

print test_mse_1, test_mse_2, test_mse_3

"""
Explanation: (a) - Deciding which is best (8 pts)
Construct separate linear models using each of the three different feature spaces. Use those models to predict responses in the test set. Compute the prediction performance of each model. Which model is the best overall? (Don't worry about statistical comparison, just choose the one with the lowest average MSE across the responses.)
You should fit a model for each feature space using ridge regression. Use cross-validation (pick your flavor) to select the best ridge parameter separately for each. It's fine here to use the same ridge parameter for each of the $m$ responses.
End of explanation
"""


# Here's a function that computes R^2 of a_hat predicting a
Rsq = lambda a, a_hat: 1 - (a - a_hat).var() / a.var()

# Here it's probably useful to define a cv_ridge function that you can call a bunch of times
def cv_ridge(x, y, x_test, y_test):
    lambdas = np.logspace(1, 7, 15) # use these lambdas

    ## YOUR CODE HERE ##
    
    return Rsq(y_test, y_test_hat)


# Now fit a model with each combination of feature spaces, and compute R^2!
# store the 7 R^2 values in all_rsqs
# use this order for the 7 models:
# 1. Just feature space 1
# 2. Just 2
# 3. Just 3
# 4. 1 & 2
# 5. 1 & 3
# 6. 2 & 3
# 7. All three (1 & 2 & 3)

all_rsqs = np.zeros(7)

## YOUR CODE HERE ##

# Now let's just look at the R^2 for each of the 7 models!
print all_rsqs

# Next, use the algebraic formulae to figure out the size of each partition
# A nice way to do this is to define a matrix V such that V.dot(all_rsqs) = partition_sizes

## YOUR CODE HERE ##

# partition_sizes should have the size of each partition (in the same units as R^2)
partition_sizes = ## YOUR CODE HERE ##

# let's compare your partition_sizes to true_variances, which is what they should be!
zip(partition_sizes, true_variances)

"""
Explanation: (b) - Variance partitioning (12 pts)
Each of these feature spaces explain something about the response. Does the variance explained by each feature space overlap with the others? Do variance partitioning to find out!
There should be 7 variance partitions: one for the unique contribution of each feature space, one for the variance explained by each pair of feature spaces, and one for the variance explained by all three spaces.
To get the sizes of these partitions, you'll need to fit models with each combination of feature spaces (each one alone, each pair, and all three together), then compute $R^2$ for each model, then use some simple algebra to compute the size of each partition.
End of explanation
"""


