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 tensorflow_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):
    xs = []
    fs = []
    xs.append(x0)
    fs.append(f(x0))
    x = xs[-1]
    for i in range(max_iter):
        xs.append(xs[-1] - step_size * df(xs[-1]))
        fs.append(f(xs[-1]))
        x = xs[-1]
        if abs(np.linalg.norm(xs[-1]) - np.linalg.norm(xs[-2])) < eps:
            break

    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 tensorflow_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/planner_v2/'
    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=list(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=list(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, list(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, list(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-')


print('Loaded learnThu8.py')

# testing tGD
# tGD([0.01, 0.02, 0.03,0.1,0.2,0.3])

# testing tGD2
# for x0 in [-4, -3, -2, -1, 0, 1, 2, 3, 4]:
#     tGD2(x0)

# test t(6)

t9()
plt.show(block=True)

