from __future__ import division
import numpy as np
from matplotlib import pyplot as plt
from sklearn import linear_model as lm

# Solutions for exercise sheet 3

def ensure_mat (X):
    if (len(X.shape) == 2):
        return X
    else:
        return X[:,np.newaxis]

def mean_squared_error (v, w):
    """
    Mean squared error between v and w

    Parameters
    ----------
    v, w: {(N,), (N,1)} array_like
          Input vectors
    
    Returns
    -------
    mse : double
    """
    return np.mean((v - w)**2)

def Xpoly (X, d):
    """
    Expand input data with polynomial basis functions

    Parameters
    ----------
    X : {(N,), (N,1)} array_like
        Input vector
    d : int
        Degree of polynomial

    Returns
    -------
    Xp : (N,d+1) array_like
         Expanded input [1, X, X^2, ..., X^d]
    """
    X = ensure_mat(X)
    return np.concatenate([X**i for i in range(d+1)], axis=1)

def pseudo_inv (X):
    return np.dot(np.linalg.inv(np.dot(X.T, X)), X.T)

def trainPoly (X, Y, d):
    """
    Train polynomial of degree d

    Parameters
    ----------
    X : {(N,), (N,1)} array_like
        Input vector
    Y : {(N,), (N,1)} array_like
        Target outputs
    d : Integer
        Degree of polynomial

    Returns
    -------
    w : (d+1,1) array_like
        Weights of trained model
    """
    Xp = Xpoly(X,d)

    try:
        # w = np.dot(pseudo_inv(Xp), Y)
        w, _, _, _ = np.linalg.lstsq(Xp, Y)
    except np.linalg.LinAlgError:
        raise TrainError('Model failed to converge')

    return w

def predictPoly (X, w):
    """
    Predictions of polynomial model

    Parameters
    ----------
    
    """
    d = len(w)-1
    Xp = Xpoly(X, d)
    Yhat = np.dot(Xp, w)
    return Yhat

def load_data ():
    # Load data sets
    trainXY = np.genfromtxt('trainXY.csv', delimiter=',')
    testXY  = np.genfromtxt('testXY.csv',  delimiter=',')

    trainX = ensure_mat(trainXY[:,0])
    trainY = ensure_mat(trainXY[:,1])
    testX  = ensure_mat(testXY[:,0])
    testY  = ensure_mat(testXY[:,1])

    return trainX, trainY, testX, testY

def exercise_3_1 ():
    trainX, trainY, testX, testY = load_data()
    
    trainMSE = []
    testMSE  = []
    degree   = []
    for d in range(10):
        degree.append(d)
        w = trainPoly(trainX, trainY, d)
        trainMSE.append(mean_squared_error(trainY, predictPoly(trainX, w)))

        testMSE.append(mean_squared_error(testY, predictPoly(testX, w)))

    plt.plot(degree, trainMSE, 'b-', degree, testMSE, 'r-')
    plt.xlabel('Model degree')
    plt.ylabel('MSE')
    plt.ylim(0,2)
    plt.title('Overfitting')

def exercise_3_2 (d):
    # generate data sets
    def f (X):
        return np.sin(2. * np.pi * X)

    def rand_data (N):
        X = np.concatenate([[0,1],np.random.uniform(0,1,size=N-2)])
        Y = f(X) + 0.2*np.random.normal(size=N)
        return ensure_mat(X), ensure_mat(Y)

    L = 100
    N = 25
    ws = []
    Xs = []
    Ys = []
    for l in range(L):
        X,Y = rand_data(N)
        w = trainPoly(X, Y, d)

        Xs.append(X)
        Ys.append(Y)
        ws.append(w)

    # Define the mean function
    def y_bar (X):
        return np.mean([predictPoly(X, w) for w in ws], axis=0)

    bias2 = np.mean([mean_squared_error(f(X), y_bar(X)) for X,Y in zip(Xs, Ys)])
    variance = np.mean([mean_squared_error(y_bar(X), predictPoly(X, w)) for X,w in zip(Xs,ws)])
    return bias2, variance, y_bar, Xs, Ys, ws