# Cell 1 from notebook

import numpy as np

def grad_U(Ui, Yij, Vj, ai, bj, reg, eta):
    """
    Takes as input Ui (the ith row of U), a training point Yij, the column
    vector Vj (jth column of V^T), reg (the regularization parameter lambda),
    and eta (the learning rate).

    Returns the gradient of the regularized loss function with
    respect to Ui multiplied by eta.
    """
    
    grad_ui = reg * Ui - (Yij - np.dot(Ui, Vj) - ai - bj) * Vj
    return grad_ui * eta

def grad_V(Vj, Yij, Ui, ai, bj, reg, eta):
    """
    Takes as input the column vector Vj (jth column of V^T), a training point Yij,
    Ui (the ith row of U), reg (the regularization parameter lambda),
    and eta (the learning rate).

    Returns the gradient of the regularized loss function with
    respect to Vj multiplied by eta.
    """
    grad_vj = reg * Vj - (Yij - np.dot(Ui, Vj) - ai - bj) * Ui
    return grad_vj * eta

def get_err(U, V, Y, a, b, reg=0.0):
    """
    Takes as input a matrix Y of triples (i, j, Y_ij) where i is the index of a user,
    j is the index of a movie, and Y_ij is user i's rating of movie j and
    user/movie matrices U and V.

    Returns the mean regularized squared-error of predictions made by
    estimating Y_{ij} as the dot product of the ith row of U and the jth column of V^T.
    """
    U_Fro_2 = 0
    for i in range(U.shape[0]):
        for j in range(U.shape[1]):
            val = U[i, j]
            U_Fro_2 += val ** 2
        
    V_Fro_2 = 0
    for i in range(V.shape[0]):
        for j in range(V.shape[1]):
            val = V[i, j]
            V_Fro_2 += val ** 2
        
    grad_sum = 0
    for i in range(Y.shape[0]):
        grad_sum += (Y[i, 2] - np.dot(U[Y[i, 0] - 1, :], V[Y[i, 1] - 1, :]) - a[Y[i, 0] - 1] - b[Y[i, 1] - 1]) ** 2
            
    err = reg / 2 * (U_Fro_2 + V_Fro_2) + grad_sum / 2
    
    return err / len(Y)


def train_model(M, N, K, eta, reg, Y, eps=0.0001, max_epochs=300):
    """
    Given a training data matrix Y containing rows (i, j, Y_ij)
    where Y_ij is user i's rating on movie j, learns an
    M x K matrix U and N x K matrix V such that rating Y_ij is approximated
    by (UV^T)_ij.

    Uses a learning rate of <eta> and regularization of <reg>. Stops after
    <max_epochs> epochs, or once the magnitude of the decrease in regularized
    MSE between epochs is smaller than a fraction <eps> of the decrease in
    MSE after the first epoch.

    Returns a tuple (U, V, err) consisting of U, V, and the unregularized MSE
    of the model.
    """ 
    
    a = np.zeros(M)
    acount = np.zeros(M)
    b = np.zeros(N)
    bcount = np.zeros(N)
    
    for idx in range(Y.shape[0]):
        i, j, Y_ij = Y[idx]
        i -= 1
        j -= 1
        a[i] += Y_ij
        acount[i] += 1
        b[j] = Y_ij
        bcount[j] += 1
        
    for idx in range(M):
        if (acount[idx] != 0):
            a[idx] /= acount[idx]
        
    for idx in range(N):
        if (bcount[idx] != 0):
            b[idx] /= bcount[idx]
    
    U = np.random.rand(M, K) - 0.5
    V = np.random.rand(N, K) - 0.5
    oldloss = get_err(U, V, Y, a, b, reg)
        
    for e in range(max_epochs):
        idx = np.random.permutation(Y.shape[0])
        for index in idx:
            i = Y[index, 0] - 1
            j = Y[index, 1] - 1
            Y_ij = Y[index, 2]
            U[i, :] -= grad_U(U[i, :], Y_ij, V[j, :], a[i], b[j], reg, eta)
            V[j, :] -= grad_V(V[j, :], Y_ij, U[i, :], a[i], b[j], reg, eta)
            
        new_loss = get_err(U, V, Y, a, b, reg)
        if (e == 0):
            loss2 = new_loss
        elif ((loss - new_loss) / (oldloss - loss2) < eps):
            loss = new_loss
            break
        loss = new_loss
    
    print(e)
    return (U, V, a, b, loss)
    
# Cell 2 from notebook

import numpy as np
import matplotlib.pyplot as plt
import pandas
import csv

with open("data/movies.txt","r", encoding="windows-1250") as f:
    reader = csv.reader(f, delimiter="\t")
    
    # Put the contents in a matrix
    movies = np.array(list(reader))

Y_train = np.loadtxt('data/train.txt').astype(int)
Y_test = np.loadtxt('data/test.txt').astype(int)
data = np.loadtxt('data/data.txt').astype(int)

M = max(max(Y_train[:,0]), max(Y_test[:,0])).astype(int) # users
N = max(max(Y_train[:,1]), max(Y_test[:,1])).astype(int) # movies
print("Factorizing with ", M, " users, ", N, " movies.")
Ks = [10,20,30,50,100]

reg = 0.0
eta = 0.03 # learning rate
E_in = []
E_out = []

# Use to compute Ein and Eout
for K in Ks:
    print(K)
    U,V, a, b, err = train_model(M, N, K, eta, reg, Y_train)
    E_in.append(err)
    E_out.append(get_err(U, V, Y_test, a, b))
    
plt.figure()
plt.plot(Ks, E_in, label='$E_{in}$')
plt.plot(Ks, E_out, label='$E_{out}$')
plt.title('Error vs. K')
plt.xlabel('K')
plt.ylabel('Error')
plt.legend()
plt.savefig('2d.png')
plt.show()

# Cell 3 from notebook

Y_train = np.loadtxt('data/train.txt').astype(int)
Y_test = np.loadtxt('data/test.txt').astype(int)

M = max(max(Y_train[:,0]), max(Y_test[:,0])).astype(int) # users
N = max(max(Y_train[:,1]), max(Y_test[:,1])).astype(int) # movies
Ks = [10,20,30,50,100]

regs = [10**-4, 10**-3, 10**-2, 10**-1, 1]
eta = 0.03 # learning rate
E_ins = []
E_outs = []

# Use to compute Ein and Eout
for reg in regs:
    E_ins_for_lambda = []
    E_outs_for_lambda = []

    for k in Ks:
        print("Training model with M = %s, N = %s, k = %s, eta = %s, reg = %s"%(M, N, k, eta, reg))
        U,V, a, b, e_in = train_model(M, N, k, eta, reg, Y_train)
        E_ins_for_lambda.append(e_in)
        eout = get_err(U, V, Y_test, a, b)
        E_outs_for_lambda.append(eout)

    E_ins.append(E_ins_for_lambda)
    E_outs.append(E_outs_for_lambda)


# Plot values of E_in across k for each value of lambda
for i in range(len(regs)):
    plt.plot(Ks, E_ins[i], label='$E_{in}, \lambda=$'+str(regs[i]))
plt.title('$E_{in}$ vs. K')
plt.xlabel('K')
plt.ylabel('Error')
plt.legend()
plt.savefig('2e_ein.png')
plt.clf()

# Plot values of E_out across k for each value of lambda
for i in range(len(regs)):
    plt.plot(Ks, E_outs[i], label='$E_{out}, \lambda=$'+str(regs[i]))
plt.title('$E_{out}$ vs. K')
plt.xlabel('K')
plt.ylabel('Error')
plt.legend()
plt.savefig('2e_eout.png')
