import pdb
import numpy as np
import matplotlib.pyplot as plt
import math

######################################################################
#
# Learning algorithms and helpers
#
######################################################################

# Ordinary least squares regression
# X is an n x d matrix
# Y is an n x 1 matrix
# returns weights: d x 1 matrix
def ols(X, y):
    return (X.T * X).I * X.T * y

# "Ridge" regression with lambda = l; otherwise like ols
# Penalizes the offset weight, which is sometimes undesirable
def olsr(X, y, l = 0):
    d = X.shape[1]
    return (X.T * X + l * np.identity(d)).I * X.T * y

# Given an order for a polynomial feature space, d
# return a *function* that maps a single value into a vector of d+1 features
def polynomialFeatures(d):
    def p(xi):
        return np.matrix([xi[0,0]**i for i in range(d+1)])
    return p

# Given X, an array of training examples, and b, a bandwidth,
# return a *function* that maps a single value into a vector of features
def RBFs(X, b):
    def g(x1, x2, b):
        return np.exp(-((np.linalg.norm(x1 - x2) / b)**2))
    return lambda xi: np.matrix([g(xi, x, b) for x in X] + [1.0])

# Given a function, such as the one generated by polynomialFeatures,
# which takes a d-dimensional vector and returns a D-dimensional
# vector, and an n x d matrix X, return an n x D matrix of feature
# vectors
def applyFeatureFun(phi, X):
    return np.vstack([phi(X[i,:]) for i in range(X.shape[0])])

# Given a D x 1 weight vector and a "feature function" that takes a
# d-dimensional vector into a D x 1 vector, return a **function** that
# takes a d-dimensional vector into a scalar regression value, but
# mapping the input into the high-dimensional feature space and then
# taking the dot product with the weights.
def makeRegressor(w, phi):
    def r(x):
        return (phi(x)*w)[0,0]
    return r

# Given an n x d matrix X of input values and an n x 1 matrix y of
# target values, and a function that maps a row of X into a
# prediction, return the root mean squared error of the predictor
# applied to the X values.
def rmse(X, y, predictor):
    return np.sqrt(sse(X, y, predictor)/len(y))

# As for rmse, but just return the sum squared error.
def sse(X, y, predictor):
    p = np.matrix([predictor(X[i,:]) for i in range(X.shape[0])]).T
    assert p.shape == y.shape
    return np.sum(np.square(p - y))

# Generic gradient descent:
# - f is a function mapping a vector of parameters to a score
#     (not strictly necessary)
# - df is a function mapping a vector of parameters to the gradient at
#     that point
# - x0 is an initial parameter vector
# - step_size is the step size
#
#
# - terminates after max_iter iterations, or when parameters have not changed
#    by more than eps from one iteration to the next
#   in every iteration updates the current param vector by the gradient
# Returns: the final parameter vector, a list of scores (one per iteration)
#   and a list of parameter vectors (one per iteration)
def gd(f, df, x0, step_size = .01, max_iter = 1000, eps = .00001):
    diff = 1
    iterCount = 0
    fs = []
    xs = []

    fs.append(f(x0))
    xs.append(x0)
    while (diff > eps and iterCount < max_iter):
        x = x0 - step_size * df(x0)
        diff = np.linalg.norm(x - x0)
        fs.append(f(x))
        xs.append(x0)

        iterCount += 1
        x0 = x
    # your code here

    return x, fs, xs

# Linear regression using gradient descent.
# X and y are as for ordinary least squares;
# w0 is the initial weight vector
# max_iter is as for gd
# l currently unused, but can be used to do ridge regression
# Returns: final weight vector, a list of scores, list of weight vectors

def gdLinReg(X, y, l = 0, step_size = 0.01, w0 = None, max_iter = 1000):
    # w is d by 1; X is n by d
    # return result is d by 1
    n = float(len(y))
    def df(w):
        return 2 * np.sum(np.multiply(X*w - y, X), axis = 0).T + l * w
    def f(w):
        return np.sum(np.square(X*w - y)) + l * float(w.T * w)

    if w0 is None: w0 = np.matrix(np.ones(X.shape[1])).T
    return gd(f, df, w0, step_size = step_size, max_iter = max_iter)

# Stochastic gradient descent
# X is an n x d matrix of input examples
# y is an n x 1 matrix of outputs (we could combine X and Y if we wanted to)
# f is a function taking a parameter vector, a 1 x d input example,
#     and a 1 x 1 output value, and returning an error; we want to
#     minimize this error in expectation over the whole data set
# df has the same inputs as f, but returns a gradient vector of the same
#     dimension as w0
# if earlyTermination is true, then terminate if params move by less
# than epsilon; otherwise, terminate after max_iter iterations.

# Returns last parameter vector, list of scores, list of param
# vectors, the iteration on which termination happened, and a Boolean
# indicating whether termination was early or not.

def sgd(X, y, f, df, w0, step_size = .01, max_iter = 20000, eps = .0000001,
            earlyTermination = False):
    n = y.shape[0]
    prev_w = w0
    fs = []; ws = [prev_w]
    for i in range(max_iter):
        j = np.random.randint(n)
        Xj = X[j]; yj = y[j]
        w = prev_w - step_size * df(prev_w, Xj, yj)
        if earlyTermination and np.all(abs(w - prev_w) < eps):
            return w, fs, ws, i, False
        fs.append(f(prev_w, Xj, yj)); ws.append(w)
        prev_w = w
    return w, fs, ws, i, True

# Linear regression computed by stochastic gradient descent.
# Parameters are as for gdLinReg.

# Returns last parameter vector, list of scores, list of param
# vectors, the iteration on which termination happened, and a Boolean
# indicating whether termination was early or not.

def sgdLinReg(X, y, l = 0, step_size = 0.01, w0 = None, max_iter = 20000):
    def df(w, Xj, yj):
        return 2 * ((Xj * w - yj) * Xj).T

    def f(w, Xj, yj):
        return np.square(Xj * w - yj)

    if w0 is None: w0 = np.matrix(np.ones(X.shape[1])).T
    return sgd(X, y, f, df, w0, step_size = step_size, max_iter = max_iter)


######################################################################
#
# Data
#
######################################################################

# Uses data set from Chapter 1 of Bishop, stored in file
# "curvefitting.txt".  They are 10 points drawn from (x, sin(2 pi x))
# with noise added (but I'm not sure how much.)

# If random is not False, it should be an integer, and instead of
# returning data from the file, we will generate a new random data set
# of that size, with 0 mean, 0.2 stdev Gaussian noise.

# if addOnes is true, return: n x 1 matrix X, n x 2 matrix F (with
# column of 1's added) and n x 1 matrix Y.

def getCurveData(addOnes = False, random = False):
    if random:
        X = np.matrix([[i / float(random)] for i in range(random + 1)])
        noise = np.random.normal(scale = 0.2, size = (random+1, 1))
        y = np.matrix([[np.sin(2 * np.pi * X[i,0])] for i in range(X.shape[0])]) + noise
    else:
        data = np.loadtxt('curvefitting.txt')
        X, y = np.matrix(data[0]).T, np.matrix(data[1]).T
    if addOnes:
        F = np.append(np.ones_like(X), X, 1)
        return X, F, y
    else:
        return X, y

# Get the Blog data

def getBlogData(n):
    dir = '/home/binder/entwurf6/teaching/2017/ai_2017/wek1/code/'
    X = np.matrix(np.genfromtxt(dir+'x_train.csv',
                                    delimiter = ','))
    if n == None:
        n = X.shape[0]
    else:
        X = X[:n]
    
    y = np.matrix(np.genfromtxt(dir+'y_train.csv', 
                                    delimiter = ',')).T[:n]
    XTest = np.matrix(np.genfromtxt(dir+'x_test.csv',
                                    delimiter = ','))
    yTest = np.matrix(np.genfromtxt(dir+'y_test.csv', 
                                    delimiter = ',')).T
    return X, y, XTest, yTest
    


######################################################################
#
# Plotting stuff
#
######################################################################

def tidyPlot(xmin, xmax, ymin, ymax, center = False, title = None,
                 xlabel = None, ylabel = None):
    plt.ion()
    plt.figure(facecolor="white")
    ax = plt.subplot()
    if center:
        ax.spines['left'].set_position('zero')
        ax.spines['right'].set_color('none')
        ax.spines['bottom'].set_position('zero')
        ax.spines['top'].set_color('none')
        ax.spines['left'].set_smart_bounds(True)
        ax.spines['bottom'].set_smart_bounds(True)
        ax.xaxis.set_ticks_position('bottom')
        ax.yaxis.set_ticks_position('left')
    else:
        ax.spines["top"].set_visible(False)    
        ax.spines["right"].set_visible(False)    
        ax.get_xaxis().tick_bottom()  
        ax.get_yaxis().tick_left()
    eps = .05
    plt.xlim(xmin-eps, xmax+eps)
    plt.ylim(ymin-eps, ymax+eps)
    if title: ax.set_title(title)
    if xlabel: ax.set_xlabel(xlabel)
    if ylabel: ax.set_ylabel(ylabel)
    return ax

def plotData(ax, x, y, style = 'ro', label = None):
    if style is None:
        ax.plot(x, y, label = label)
    else:
        ax.plot(x, y, style, label = label)
    plt.show()

# w is a (1 x 2) matrix
def plotLine(ax, w, xmin, xmax, nPts = 100):
    b = float(w[0])
    m = float(w[1])
    plotFun(ax, lambda x: m*x + b, xmin, xmax, nPts)

def plotFun(ax, f, xmin, xmax, nPts = 100, label = None):
    x = np.linspace(xmin, xmax, nPts)
    y = np.vstack([f(np.matrix([[xi]])) for xi in x])
    ax.plot(x, y, label = label)
    plt.show()


def smooth(n, vals):
    # Run a box filter of size n
    x = sum(vals[0:n])
    result = [x]
    for i in range(n, len(vals)):
        x = x - vals[i-n] + vals[i]
        result.append(x)
    return result

######################################################################
#
# Tests
#
######################################################################

######################################################################
#
# Ordinary least squares in 2D (one constant input dimension)

def t1():
    X, F, y = getCurveData(True)

    w = ols(F, y)

    print ('w', w.T)
    xmin, xmax = float(min(X)), float(max(X))
    ymin, ymax = float(min(y)), float(max(y))
    ax = tidyPlot(xmin, xmax, ymin, ymax, xlabel = 'x', ylabel = 'y')
    plotData(ax, X, y)
    plotLine(ax, w, xmin, xmax)

######################################################################
#
# Ordinary least squares in polynomial feature spaces
# Plots predictors

def t2(ds = range(1, 10)):
    X, y = getCurveData()
    xmin, xmax = float(min(X)), float(max(X))
    ymin, ymax = float(min(y)), float(max(y))
    ax = tidyPlot(xmin, xmax, ymin-1, ymax+1, xlabel = 'x', ylabel = 'y')

    plotData(ax, X, y)

    for d in ds:
        phi = polynomialFeatures(d)
        phiD = applyFeatureFun(phi, X)
        w = ols(phiD, y)
        predictor = makeRegressor(w, phi)

        plotFun(ax, predictor, xmin, xmax, label = str(d))
        print ('Order', d, 'Training RMSE', rmse(X, y, predictor))
        print ('     w', w.T)
    ax.legend(loc="upper left", bbox_to_anchor=(1,1))

######################################################################
#
# Ordinary least squares in polynomial feature spaces
# Plots train and test error versus order of polynomial basis

def t3(ds = range(1, 10)):
    X, y = getCurveData()
    XTest, yTest = getCurveData(random = 10)
    xmin, xmax = float(min(X)), float(max(X))
    ymin, ymax = float(min(y)), float(max(y))
    ax = tidyPlot(xmin, xmax, ymin-1, ymax+1, xlabel = 'x', ylabel = 'y')
    plotData(ax, X, y)
    plotData(ax, XTest, yTest, style = 'go')
    trainErr = []
    testErr = []

    for d in ds:
        phi = polynomialFeatures(d)
        phiD = applyFeatureFun(phi, X)
        w = ols(phiD, y)
        predictor = makeRegressor(w, phi)
        plotFun(ax, predictor, xmin, xmax)
        trainErr.append(rmse(X, y, predictor))
        testErr.append(rmse(XTest, yTest, predictor))
        print ('Order', d, 'Training RMSE', trainErr[-1])
        print ('        Test RMSE', testErr[-1])

    ax = tidyPlot(0, len(ds), 0, max(max(testErr), max(trainErr)), center =True,
                      xlabel = 'Polynomial order', ylabel = 'RMSE')
    ax.plot(ds, trainErr, label = 'train err')
    ax.plot(ds, testErr, label = 'test err')
    ax.legend()

######################################################################
#
# Ordinary least squares in RBF feature spaces
# Plots train and test error versus bandwidth

# LPK: these are good b values (.01, .05, .1, .5, 1, 2, 4)

def tRBF(bs = []):
    X, y = getCurveData()
    XTest, yTest = getCurveData(random = 10)
    xmin, xmax = float(min(X)), float(max(X))
    ymin, ymax = float(min(y)), float(max(y))
    ax = tidyPlot(xmin, xmax, ymin-1, ymax+1, xlabel = 'x', ylabel = 'y')
    plotData(ax, X, y)
    trainErr = []
    testErr = []

    for b in bs:
        phi = RBFs(X, b)
        phiD = applyFeatureFun(phi, X)
        w = olsr(phiD, y, 1e-5) 
        predictor = makeRegressor(w, phi)
        plotFun(ax, predictor, xmin, xmax, label = str(b))
        trainErr.append(rmse(X, y, predictor))
        testErr.append(rmse(XTest, yTest, predictor))
        print ('Bandwidth', b, 'Training RMSE', trainErr[-1])
        print ('        Test RMSE', testErr[-1])
        print ('    w', w.T)
    ax.legend(loc = 'best')

    if len(bs) > 1:
        ax = tidyPlot(min(bs), max(bs), 0, max(max(testErr), max(trainErr)),
                        xlabel = 'Bandwidth', ylabel = 'RMSE')
        ax.plot(bs, trainErr, label = 'train err')
        ax.plot(bs, testErr, label = 'test err')
        ax.legend(loc = 'best')


######################################################################
#
# Ordinary least squares in polynomial feature spaces
# Plots train and test error versus size of training set

def t4(trainSizes = (10, 15, 20, 30, 50, 100), showData = False):
    XTest, yTest = getCurveData(random = 100)
    xmin, xmax = float(np.min(XTest)), float(np.max(XTest))
    ymin, ymax = float(np.min(yTest)), float(np.max(yTest))
    trainErr = []
    testErr = []
    phi = polynomialFeatures(9)

    if not showData:
        ax = tidyPlot(xmin, xmax, ymin-1, ymax+1, xlabel = 'x', ylabel = 'y',
                          title = 'Predictors for different train sizes')

    for ts in trainSizes:
        X, y = getCurveData(random = ts-1)
        if showData:
            ax = tidyPlot(xmin, xmax, ymin-1, ymax+1, xlabel = 'x', ylabel = 'y',
                          title = 'Train size = ' + str(ts))
            plotData(ax, X, y)
        phiD = applyFeatureFun(phi, X)
        w = ols(phiD, y)
        predictor = makeRegressor(w, phi)
        plotFun(ax, predictor, xmin, xmax, label=str(ts))
        trainErr.append(rmse(X, y, predictor))
        testErr.append(rmse(XTest, yTest, predictor))
        print ('Train size', ts, 'Training RMSE', trainErr[-1])
        print ('        Test RMSE', testErr[-1])
        print ('     w', w.T)
    ax.legend(loc="upper left", bbox_to_anchor=(1,1))
    if len(trainSizes) > 1:
        ax = tidyPlot(0, max(trainSizes), 0, max(testErr), center = True,
                      xlabel = 'Training set size', ylabel = 'RMSE')
        ax.plot(trainSizes, trainErr, label = 'train err')
        ax.plot(trainSizes, testErr, label = 'test err')
        ax.legend(loc="upper right")


######################################################################
#
# Ridge regression in polynomial feature spaces
# Plot train and test error versus ridge parameter lambda

def t5(logLambdaValues = (-50, -40, -30, -20, -15, -10, -1, 0, 1, 10)):
    XTest, yTest = getCurveData(random = 100)
    xmin, xmax = float(np.min(XTest)), float(np.max(XTest))
    ymin, ymax = float(np.min(yTest)), float(np.max(yTest))
    trainErr = []
    testErr = []
    phi = polynomialFeatures(9)

    ax = tidyPlot(xmin, xmax, ymin-1, ymax+1,
                      xlabel = 'x', ylabel = 'y',
                          title = 'Predictors for different lambda values')
    X, y = getCurveData()
    phiD = applyFeatureFun(phi, X)
    plotData(ax, X, y)
    for llv in logLambdaValues:
        w = olsr(phiD, y, np.exp(llv))
        predictor = makeRegressor(w, phi)
        plotFun(ax, predictor, xmin, xmax, label=str(llv))
        trainErr.append(rmse(X, y, predictor))
        testErr.append(rmse(XTest, yTest, predictor))
        print ('Log lambda', llv, 'Training RMSE', trainErr[-1])
        print ('        Test RMSE', testErr[-1])
        print (w.T)
    ax.legend(loc="upper left", bbox_to_anchor=(1,1))

    if len(logLambdaValues) > 1:
        ax = tidyPlot(min(logLambdaValues), max(logLambdaValues),
                        0, max(testErr), center = True,
                        xlabel = 'Log lambda', ylabel = 'RMSE')
        ax.plot(logLambdaValues, trainErr, label = 'train err')
        ax.plot(logLambdaValues, testErr, label = 'test err')
        ax.legend(loc="best")
    

######################################################################
#
# Batch gradient descent in 2D feature space
# Plots error as a function of iteration number for different step sizes
# Individual plots of trajectory of w during optimization

def t6(learning_rates = (.0001, .001, .01, .05, .07, .075, .1)):
    X, F, y = getCurveData(True)
    ax = tidyPlot(0, 1000, 2, 8,
                      xlabel = 'iteration', ylabel = 'err')
    for lr in learning_rates:
        w, scores, vals = gdLinReg(F, y, step_size = lr)
        
        plotData(ax, range(len(scores)), scores, style = None,
                     label = str(lr))
        print ('lr', lr, 'w', w.T, 'err', scores[-1])
        dmin = min(-1.5, np.min(vals)); dmax = max(1.5, np.max(vals))
        nax = tidyPlot(dmin, dmax, dmin, dmax,
                       xlabel = 'w0', ylabel = 'w1',
                       title = 'step size = '+str(lr), center = True)
        plotData(nax, [float(xd) for (xd, yd) in vals],
                      [float(yd) for (xd, yd) in vals],
                      style = 'bo-')
    ax.legend(loc="upper left", bbox_to_anchor=(1,1))

######################################################################
#
# Stochastic  gradient descent in 2D feature space
# Plots error as a function of iteration number for different step sizes
# Individual plots of trajectory of w during optimization

def t7(learning_rates = (.0001, .001, .01, .1, .5, 1)):
    X, F, y = getCurveData(True)
    ax = tidyPlot(0, 1000, 0, 100,
                      xlabel = 'iteration', ylabel = 'err')
    for lr in learning_rates:
        w, scores, vals, iters, failed = sgdLinReg(F, y, step_size = lr)
        
        smoothScores = smooth(len(y)*4, [float(s) for s in scores])
        plotData(ax, range(len(smoothScores)), smoothScores,
                     style = None, label = str(lr))
        print ('lr', lr, 'w', w.T, 'err', smoothScores[-1])
        dmin = min(-1.5, np.min(vals)); dmax = max(1.5, np.max(vals))
        nax = tidyPlot(dmin, dmax, dmin, dmax,
                       xlabel = 'w0', ylabel = 'w1',
                       title = 'step size = '+str(lr), center = True)
        plotData(nax, [float(xd) for (xd, yd) in vals],
                      [float(yd) for (xd, yd) in vals],
                      style = 'bo-')
    ax.legend(loc="best")

######################################################################
#
# Batch gradient descent in polynomial feature space
# Plots train and test error as a function of iteration number

def t8(order = 9):
    goodw = np.matrix(\
      [[  3.49512436e-01,   2.32326298e+02,  -5.32086773e+03,   4.85596769e+04,
         -2.31598783e+05,   6.39931981e+05,  -1.06162004e+06,   1.04222515e+06,
         -5.57590377e+05,   1.25180846e+05]]).T
    zeros = np.matrix(np.zeros([order+1, 1]))
    medw = np.matrix(\
      [[  3.0e-01,   2.0e+02,  -5.0e+03,   5.0e+04,
         -2.0e+05,   6.0e+05,  -1.0e+06,   1.0e+06,
         -5.0e+05,   1.0e+05]]).T

    X, y = getCurveData()
    xmin, xmax = float(min(X)), float(max(X))
    ymin, ymax = float(min(y)), float(max(y))
    XTest, yTest = getCurveData(random = 100)

    phi = polynomialFeatures(order)
    phiD = applyFeatureFun(phi, X)
    lr = 0.005
    w, scores, vals = gdLinReg(phiD, y, step_size = .01, w0 = zeros,
                                   max_iter = 10000)
    print ('w', w.T, 's', scores[-1], 'iter', len(scores))
    iVals = []; trainVals = []; testVals = []
    axPred = tidyPlot(xmin, xmax, ymin-1, ymax+1, xlabel = 'x', ylabel = 'y')
    for i in range(0, len(vals), 100):
        predictor = makeRegressor(vals[i], phi)
        plotFun(axPred, predictor, xmin, xmax, label = 'iters = '+str(i))
        iVals.append(i)
        trainVals.append(rmse(X, y, predictor))
        testVals.append(rmse(XTest, yTest, predictor))
    plotData(axPred, X, y)
    axPred.legend(loc="upper left", bbox_to_anchor=(1,1))
        
    ax = tidyPlot(0, np.max(iVals), 0, 1, #np.max(testVals),
                   xlabel = 'iteration', ylabel = 'rmse')
    plotData(ax, iVals, trainVals, label = 'training error', style = 'bo-')
    plotData(ax, iVals, testVals, label = 'testing error', style = 'ro-')
    ax.legend(loc="best")

def t9(order = 9):
    zeros = np.matrix(np.zeros([order+1, 1]))
    goodw = np.matrix(\
      [[  3.49512436e-01,   2.32326298e+02,  -5.32086773e+03,   4.85596769e+04,
         -2.31598783e+05,   6.39931981e+05,  -1.06162004e+06,   1.04222515e+06,
         -5.57590377e+05,   1.25180846e+05]]).T
    medw = np.matrix(\
      [[  3.0e-01,   2.0e+02,  -5.0e+03,   5.0e+04,
         -2.0e+05,   6.0e+05,  -1.0e+06,   1.0e+06,
         -5.0e+05,   1.0e+05]]).T
    randw = np.matrix(np.random.randn(order+1,1))

    X, y = getCurveData()
    xmin, xmax = float(min(X)), float(max(X))
    ymin, ymax = float(min(y)), float(max(y))
    XTest, yTest = getCurveData(random = 100)

    phi = polynomialFeatures(order)
    phiD = applyFeatureFun(phi, X)
    lr = 0.05
    w, scores, vals, iters, failed = sgdLinReg(phiD, y, step_size = lr,
                                                   max_iter = 100000,
                                                   w0 = randw)
    print ('w', w.T, 's', scores[-1], 'iter', len(scores))
    iVals = []; trainVals = []; testVals = []
    axPred = tidyPlot(xmin, xmax, ymin-1, ymax+1, xlabel = 'x', ylabel = 'y')
    for i in range(0, len(vals), 5000):
        predictor = makeRegressor(vals[i], phi)
        plotFun(axPred, predictor, xmin, xmax, label = 'iters = '+str(i))
        iVals.append(i)
        trainVals.append(rmse(X, y, predictor))
        testVals.append(rmse(XTest, yTest, predictor))
    plotData(axPred, X, y)
        
    ax = tidyPlot(0, np.max(iVals), 0, 1, #np.max(testVals),
                   xlabel = 'iteration', ylabel = 'rmse')
    plotData(ax, iVals, trainVals, label = 'training error', style = 'bo-')
    plotData(ax, iVals, testVals, label = 'testing error', style = 'ro-')
    ax.legend(loc="upper left", bbox_to_anchor=(1,1))


# Blog feedback data, with analytic ridge regression     
def t10(n = None,
        logLambdaValues = [-14, -12, -10, -8, -6, -4, -2, 0, 2, 4, 5, 6, 7,
                               8, 9, 10, 12, 14, 16, 18, 20]):
    X, y, XTest, yTest = getBlogData(n = n)
    print (X.shape, y.shape)
    iVals = []; trainVals = []; testVals = []
    for llv in logLambdaValues:
        w = olsr(X, y, np.exp(llv))
        predictor = lambda x: x * w
        iVals.append(llv)
        trainVals.append(rmse(X, y, predictor))
        testVals.append(rmse(XTest, yTest, predictor))
        print (llv, trainVals[-1], testVals[-1])
        
    ax = tidyPlot(np.min(iVals), np.max(iVals), 28, 40, #np.max(testVals),
                   xlabel = 'Log lambda', ylabel = 'rmse',
                   title = 'Training data size = '+str(n))
    plotData(ax, iVals, trainVals, label = 'training error', style = 'bo-')
    plotData(ax, iVals, testVals, label = 'testing error', style = 'ro-')
    ax.legend(loc="upper left", bbox_to_anchor=(1,1))

# Blog feedback data, with gradient regression.
# Add ridge!

# lambda = 80; rmse = 31.8; mse = 1024
# lambda = ? ; mse = 1015   (seems untrustworthy)
# log lambda = 2.8
# lambda 100, MSE 894

# lr = .000000005 good for whole data set (or divide by 5)

def t11(lr = .000000005, n = None): 
    X, y, XTest, yTest = getBlogData(n = n)
    print (X.shape, y.shape)
    iVals = []; trainVals = []; testVals = []
    #w, scores, vals = gdLinReg(X, y, step_size = lr, max_iter = 20000)
    w, scores, vals, iters, failed = \
                            sgdLinReg(X, y, 
                                          step_size = lr, max_iter = 2000000)
    print ('s', scores[-1], 'iter', len(scores))
    iVals = []; trainVals = []; testVals = []
    for i in range(0, len(vals), 50000):
        predictor = lambda x: float(x * vals[i])
        iVals.append(i)
        trainVals.append(rmse(X, y, predictor))
        testVals.append(rmse(XTest, yTest, predictor))

    print (trainVals)
    print (testVals)
        
    ax = tidyPlot(0, np.max(iVals), 0, 400, #np.max(testVals),
                   xlabel = 'iteration', ylabel = 'rmse')
    plotData(ax, iVals, trainVals, label = 'training error', style = 'bo-')
    plotData(ax, iVals, testVals, label = 'testing error', style = 'ro-')
    ax.legend(loc="upper left", bbox_to_anchor=(1,1))

def tGD(step_sizes = []):
    def f(x):
        return (2 * x + 3)**2
    def df(x):
        return 2 * 2 * (2 * x + 3)
    x0 = 0
    for ss in step_sizes:
        x, fs, xs = gd(f, df, x0, step_size = ss)
        print ('ss', ss, 'x', x)
        nax = tidyPlot(-4, 1, 0, 15,
                       xlabel = 'x', ylabel = 'f(x)',
                       title = 'step size = '+str(ss), center = True)
        plotFun(nax, f, -4, 1)
        plotData(nax, xs, fs, style = 'ro-')

def tGD2(x0 = 0):
    def f(x):
        return (x - 2) * (x - 3) * (x + 3) * (x + 1)
    def df(x):
        return 9 - (22 * x) - (3 * x**2) + (4 * x**3)
    ss = 0.01
    x, fs, xs = gd(f, df, x0, step_size = ss)
    nax = tidyPlot(-4, 4, -25, 25,
                       xlabel = 'x', ylabel = 'f(x)',
                       title = 'step size = '+str(ss), center = True)
    plotFun(nax, f, -4, 4)
    plotData(nax, xs, fs, style = 'ro-')
        
if __name__ == '__main__':
    print ('Loaded learn.py')
    # tGD2()
    stepSizes = [0.01]
    tGD(stepSizes)
    plt.show(block=True)

