# -*- coding: utf-8 -*-
#!/usr/bin/python

'''
Adapted from example code by Albert Au Yeung (2010)
http://www.quuxlabs.com/blog/2010/09/matrix-factorization-a-simple-tutorial-and-implementation-in-python/#source-code

An implementation of matrix factorization...

@INPUT:
    R     : a matrix to be factorized, dimension N x M
    P     : an initial matrix of dimension N x K
    Q     : an initial matrix of dimension M x K
    K     : the number of latent features
    steps : the maximum number of steps to perform the optimisation
    alpha : the learning rate
    beta  : the regularization parameter
@OUTPUT:
    the final matrices P and Q, steps taken and error

'''
import numpy 


def matrix_factorization(R, P, Q, K, steps=5000, alpha=0.0002, beta=0.02):
    half_beta = beta / 2.0
    for step in xrange(steps):
        for i in xrange(len(R)):
            for j in xrange(len(R[i])):
                if R[i,j] > 0:
                    eij = R[i,j] - numpy.dot(P[i,:],Q[j,:])
                    for k in xrange(K):
                        P[i,k] += alpha * (2 * eij * Q[j,k] - beta * P[i,k])
                        Q[j,k] += alpha * (2 * eij * P[i,k] - beta * Q[j,k])
        e = 0
        for i in xrange(len(R)):
            for j in xrange(len(R[i])):
                if R[i,j] > 0:
                    e += (R[i,j] - numpy.dot(P[i,:],Q[j,:]))**2
                    for k in xrange(K):
                        e += half_beta * (P[i,k]*P[i,k] + Q[j,k]*Q[j,k])
        if e < 0.001:
            break
    return P, Q, step, e