import numpy as np
import scipy.linalg
from scipy.optimize import minimize
from math import sqrt
import time

data = np.loadtxt('CASP.csv', delimiter=',', skiprows=1)

y = data[:, 0]
X = data[:, 1:]


def split_train_test(X, y, fraction_train=9.0 / 10.0):
    end_train = round(X.shape[0] * fraction_train)
    X_train = X[0:end_train, ]
    y_train = y[0:end_train]
    X_test = X[end_train:, ]
    y_test = y[end_train:]
    return X_train, y_train, X_test, y_test


def normalize_features(X_train, X_test):
    mean_X_train = np.mean(X_train, 0)
    std_X_train = np.std(X_train, 0)
    std_X_train[std_X_train == 0] = 1
    X_train_normalized = (X_train - mean_X_train) / std_X_train
    X_test_normalized = (X_test - mean_X_train) / std_X_train
    return X_train_normalized, X_test_normalized


X_train, y_train, X_test, y_test = split_train_test(X, y)
X_train, X_test = normalize_features(X_train, X_test)
X_train = np.concatenate((X_train, np.ones((X_train.shape[0], 1))), 1)
X_test = np.concatenate((X_test, np.ones((X_test.shape[0], 1))), 1)


def map_estimate(X, y, sigma=1.0, tau=10.):
    lamb = sqrt(tau) * np.identity(X.shape[1])
    X_tilde = np.concatenate((X/sqrt(sigma), lamb))
    Y_tilde = np.concatenate((y/sqrt(sigma), np.zeros(X.shape[1])))
    q, r = np.linalg.qr(X_tilde)
    return scipy.linalg.solve_triangular(r, np.dot(q.T, Y_tilde),
                                         check_finite=False)


def problem_4():
    w = map_estimate(X_train, y_train)
    predictions = np.dot(X_test, w)
    rmse = np.linalg.norm(predictions - y_test) / sqrt(len(y_test))
    print w
    print rmse


def log_posterior(w, sigma=1.0, tau=10):
    return (1 / (2 * sigma) * np.linalg.norm(y_train - np.dot(X_train, w)) ** 2
            + 1 / (2 * tau) * np.linalg.norm(w) ** 2)


def gradient(w, sigma=1.0, tau=10):
    return (- 1 / sigma * np.dot(X_train.T, y_train - np.dot(X_train, w))
            + 1 / tau * w)


def verify_gradient(eps=1e-10):
    w = np.ones(X_train.shape[1])
    for i in range(X_train.shape[1]):
        epsilon = np.zeros(X_train.shape[1])
        epsilon[i] = eps
        print (log_posterior(w + epsilon)
               - log_posterior(w - epsilon)) / (2 * eps)
        print gradient(w)[i]


def problem_5():
    x0 = np.zeros(X_train.shape[1])
    w = minimize(log_posterior, x0, jac=gradient, method='L-BFGS-B',
                 options={'maxiter': 100}).x
    predictions = np.dot(X_test, w)
    rmse = np.linalg.norm(predictions - y_test) / sqrt(len(y_test))
    print w
    print rmse


if __name__ == "__main__":
    problem_4()
    problem_5()