import sys
from itertools import islice

import autograd.numpy as np
from autograd import grad


def get_ratings(filename):
    with open(filename) as fh:
        for line in fh:
            yield map(int, line.strip().split())


def get_train_test(filename):
    l = [(i, j) for (i, j, _) in get_ratings(filename)]
    n = max(i for (i, _) in l)
    m = max(j for (_, j) in l)
    g = get_ratings(filename)
    train = islice(g, 100000)
    test = islice(g, 100000)
    return n, m, list(train), list(test)


def kl(muu, muv, sigmau, sigmav):
    sigma = 5.
    r = (np.log(sigma / sigmav).sum() + np.log(sigma / sigmau).sum()
         + np.sum(sigmau ** 2 / (2 * sigma ** 2))
         + np.sum(sigmav ** 2 / (2 * sigma ** 2))
         + np.sum(muu ** 2 / (2 * sigma ** 2))
         + np.sum(muv ** 2 / (2 * sigma ** 2))
         )
    return r


def likelihood(params, r):
    su, mu, sv, mv = params
    s = len(mu)
    ns = 50
    u = su * np.random.normal(size=(ns, s)) + mu
    v = sv * np.random.normal(size=(ns, s)) + mv
    p = np.sum(u * v, axis=1)
    m = np.sum((r - p) ** 2 / 2.)
    return m / float(ns)

gl = grad(likelihood)


def sgd(ratings, muu, muv, sigmau, sigmav):
    lamb = 0.05
    sigma = 5.
    eps = 1e-5

    for e in xrange(10):
        print l(ratings, muu, muv, sigmau, sigmav)
        # gradient update on the KL part
        sigmau = sigmau + lamb * 1. / sigmau
        sigmav = sigmav + lamb * 1. / sigmav
        sigmau = sigmau - lamb * sigmau / (sigma ** 2)
        sigmav = sigmav - lamb * sigmav / (sigma ** 2)
        muu = muu - lamb * muu / (sigma ** 2)
        muv = muv - lamb * muv / (sigma ** 2)

        for k, (i, j, r) in enumerate(ratings):
            if k % 1000 == 0:
                print k
            su, mu, sv, mv = sigmau[i], muu[i], sigmav[j], muv[j]
            gsu, gmu, gsv, gmv = gl((su, mu, sv, mv), r)
            sigmau[i] = np.maximum(sigmau[i] - lamb * gsu, eps)
            muu[i] = muu[i] - lamb * gmu
            sigmav[j] = np.maximum(sigmav[j] - lamb * gsv, eps)
            muv[j] = muv[j] - lamb * gmv


def l(ratings, muu, muv, sigmau, sigmav):
    res = kl(muu, muv, sigmau, sigmav)
    for (i, j, r) in ratings:
        su, mu, sv, mv = sigmau[i], muu[i], sigmav[j], muv[j]
        res += likelihood((su, mu, sv, mv), r)
    return res


if __name__ == "__main__":
    n, m, train, test = get_train_test(sys.argv[1])
    print len(train)
    k = 2
    muu, sigmau = 0.01 * np.random.normal(size=(n+1, k)), 0.01 * np.ones((n+1, k))
    muv, sigmav = 0.01 * np.random.normal(size=(m+1, k)), 0.01 * np.ones((n+1, k))
    sgd(train, muu, muv, sigmau, sigmav)
    print l(train, muu, muv, sigmau, sigmav)