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

######################################################################
#
# 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 an order for a polynomial feature space, order
# return a *function* that maps an input vector of dimension d to
# an output vector of dimension D by taking all combinations
def polynomialFeaturesN(order):
    def p(xi):
        d = xi.shape[1]
        # assume xi is matrix [[xi1, ...., xid]]
        features = []
        for o in range(order+1):
            indexTuples = itertools.combinations_with_replacement(range(d), o)
            for it in indexTuples:
                features.append(mul(xi[0, i] for i in it))
        return np.matrix([features])
    return p

def mul(seq):
    return reduce(operator.mul, seq, 1)


# 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

def makeLogisticRegressor(w, phi):
    def r(x):
        return s((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
# 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,
           lrDecay = 1.0):
    prev_x = x0
    prev_f = f(x0)
    fs = [prev_f]; xs = [prev_x]
    for i in range(max_iter):
        x = prev_x - step_size * df(prev_x)
        if np.all(abs(x - prev_x) < eps): return x, fs, xs
        fs.append(f(x)); xs.append(x)
        prev_x = x
        step_size *= lrDecay
    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 is 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 / n + l * w 
    def f(w):
        return np.sum(np.square(X*w - y)) / n + 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)


# Sigmoid function

def s(z):
    return 1.0 / (1 + np.exp(-z))

# Logistic regression using gradient descent.
# X is a matrix of feature vectors; y is a vector of labels in {0, 1}
# w0 is the initial weight vector
# max_iter is as for gd
# l currently unused, but can be used for regularization
# Returns: final weight vector, a list of scores, list of weight vectors

def gdLogReg(X, y, l = 0, step_size = 0.01, w0 = None, max_iter = 1000,
                 eps = .00001):
    # w is d by 1; X is n by d
    # return result is d by 1
    n = float(len(y))

    def df(w):
        return np.sum(np.multiply(s(X*w) - y, X), axis = 0).T + l * w

    def f(w):
        z = s(X*w)
        score =  - (np.sum(np.multiply(y, np.log(z)) + \
                      np.multiply(1.0 - y, np.log(1.0 - z)))) \
                  + l * float(w.T * w)
        return score

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

# 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

def superSimpleSeparable(addOnes = False):
    X = np.matrix([[2, 3],
                   [3, 2],
                   [9, 10],
                   [10, 9]])
    y = np.matrix([[1, 1, 0, 0]]).T
    if addOnes:
        X = np.append(np.ones_like(y), X, 1)
    return X, y

def superSimpleSeparable2(addOnes = False):
    X = np.matrix([[2, 5],
                   [3, 2],
                   [9, 6],
                   [12, 5]])
    y = np.matrix([[1, 0, 1, 0]]).T
    if addOnes:
        X = np.append(np.ones_like(y), X, 1)
    return X, y

def xor(addOnes = False):
    X = np.matrix([[1, 1],
                   [2, 2],
                   [1, 2],
                   [2, 1]])
    y = np.matrix([[1, 1, 0, 0]]).T
    if addOnes:
        X = np.append(np.ones_like(y), X, 1)
    return X, y

def xor_more(addOnes = False):
    X = np.matrix([[1, 1], [2, 2], [1, 2], [2, 1],
                   [2, 3], [4, 1], [1, 3], [3, 3]])

                   
    y = np.matrix([[1, 1, 0, 0, 1, 1, 0, 0]]).T
    if addOnes:
        X = np.append(np.ones_like(y), X, 1)
    return X, y

def multimodalData(modes = None, numPerMode = 20,
                        numModes = 2,
                        modeCov = np.eye(2, 2)):
    Xs = []
    Ys = []
    if modes is None: