#%% [markdown]
# #RecSym Playground
# The notebook contains several demo models that I tried when learning recommender systems from different sources.

#%% Matrix Factorization in Torch
# link to tutorial: https://www.ethanrosenthal.com/2017/06/20/matrix-factorization-in-pytorch/
import torch
from scipy.sparse import rand
import numpy as np

tasks = ['A','B','C','D','N-back']
constructs = ['A1','A2','A3','A4','C5','C6']
n_latent_vars = 2

# generate random hits
hits = rand(len(tasks), len(constructs), density=0.1, format='csr')
hits.data = (np.random.randint(1, 100, size=hits.nnz).astype(np.float64))
hits = hits.toarray()

class MatrixFactorization(torch.nn.Module):
  def __init__(self,n_tasks, n_constructs, n_factors):
    super().__init__()
    self.task_factors = torch.nn.Embedding(n_tasks, n_factors, sparse=True)
    self.construct_factors = torch.nn.Embedding(n_constructs, n_factors, sparse=True)
    self.task_biases = torch.nn.Embedding(n_tasks, 1, sparse=True)
    self.construct_biases = torch.nn.Embedding(n_constructs, 1, sparse=True)

  def forward(self, task, construct):
    pred = self.task_biases(task) + self.construct_biases(construct)
    pred += (self.task_factors(task) * self.construct_factors(construct)).sum(dim=1, keepdim=True)
    return pred.squeeze()

model = MatrixFactorization(len(tasks), len(constructs), n_latent_vars)

loss_func = torch.nn.MSELoss()

optimizer = torch.optim.SGD(model.parameters(), lr=1e-5)

# shuffle data
_tasks, _constructs = hits.nonzero()
p = np.random.permutation(len(_tasks))
_tasks, _constructs = _tasks[p], _constructs[p]

for t, c in zip(*(_tasks, _constructs)):
    hit = torch.FloatTensor([hits[t, c]])
    t = torch.LongTensor([t])
    c = torch.LongTensor([c])
    
    prediction = model(t, c)
    loss = loss_func(prediction, hit)
    loss.backward()
    optimizer.step()

display(model)
#%% TODO collaborative filtering - implicit versus explicit rating

#%% TODO item-based CF vs. user-based CF

#%% TODO matrix factorization with SVD