import sys
from itertools import groupby, chain
import autograd.numpy as np
from autograd import grad
import autograd.scipy as sp
from math import sqrt

b = np.array([-1e100, -4., -2., 2., 4., 1e100])


def get_ratings():
    with open(sys.argv[1]) as fh:
        for line in fh:
            i, j, r = map(int, line.strip().split())
            yield (j - 1, r)


def split_train_test():
    l = list(get_ratings())
    l.sort(key=lambda x: x[0])
    for k, g in groupby(l, key=lambda x: x[0]):
        g = list(g)
        n = int(0.95 * len(g))
        train = g[:n]
        test = g[n:]
        yield train, test


def build_train_test():
    train, test = zip(*split_train_test())
    train = np.array(list(chain.from_iterable(train)))
    test = np.array(list(chain.from_iterable(test)))
    return train, test


def log_posterior(ws, batch):
    w = ws[:-1]
    sigmas = ws[-1]
    indices = batch[:, 0]
    ratings = batch[:, 1]
    m = np.dot(X, w)[indices]
    n = b[ratings]
    o = b[ratings - 1]
    a1 = sp.stats.norm.logcdf(1. / sigmas * (n - m))
    a2 = sp.stats.norm.logcdf(1. / sigmas * (o - m))
    return - np.sum(np.log(1 - np.exp(a2 - a1)) + a1) + 0.5 * np.sum(w * w)


def log_posterior_bis(ws, batch):
    w = ws[:-1]
    sigmas = ws[-1]
    indices = batch[:, 0]
    ratings = batch[:, 1]
    m = np.dot(X, w)[indices]
    return - np.sum(-(m - ratings) ** 2 / (2 * sigmas)) + 0.5 * np.sum(w * w)

grad_log_posterior = grad(log_posterior)


def sgd():
    b1 = 0.9
    b2 = 0.999
    b1t = b1
    b2t = b2
    eps = 1e-8
    alpha = 0.001
    w = np.ones(X.shape[1] + 1)
    m = np.zeros(X.shape[1] + 1)
    v = np.zeros(X.shape[1] + 1)
    for i in xrange(100):
        print i
        for k in xrange(len(train) / 100):
            batch = train[k * 100: k * 100 + 100, :]
            g = grad_log_posterior(w, batch)
            m = b1 * m + (1. - b1) * g
            v = b2 * v + (1. - b2) * (g * g)
            mt = m / (1. - b1t)
            vt = v / (1. - b2t)
            tmp = alpha * mt / (np.sqrt(vt) + eps)
            w = w - tmp
        b1t *= b1
        b2t *= b2
        print sqrt(compute_mse(w, test))


def compute_mse(w, t):
    w = w[:-1]
    ratings = np.dot(X, w)
    ratings = np.searchsorted(b, ratings)
    ratings = ratings[t[:, 0]]
    s = np.sum((ratings - t[:, 1]) ** 2)
    return float(s) / len(t)


if __name__ == "__main__":
    train, test = build_train_test()
    X = np.loadtxt(sys.argv[2])
    w = np.ones(X.shape[1] + 1)
    #print grad_log_posterior(w, train)
    sgd()