import pandas as pd
import numpy as np

movies = pd.read_csv('ml-25m/movies.csv')
print(movies.shape)
movies.columns

year = 2010
movies['year'] = movies.title.str.extract(r'\((\d{4})\)').astype("float")
movie_id_removed = movies.query('year < @year').movieId.tolist()
movies = movies.query('year >= @year')
movies

len(movie_id_removed)

movies_mapping = movies[['movieId','title']].set_index('movieId').to_dict()['title']

ratings = pd.read_csv('ml-25m/ratings.csv', parse_dates=['timestamp'])
print(ratings.columns)
print(ratings.shape)
ratings.userId.nunique()

ratings = ratings.query('movieId not in @movie_id_removed')
ratings['rating'] = 1
ratings['timestamp'] = pd.to_datetime(ratings['timestamp'], unit='s')
ratings

from sklearn.preprocessing import LabelEncoder

user_encoder = LabelEncoder()
movie_encoder = LabelEncoder()
ratings['userId'] = user_encoder.fit_transform(ratings.userId)
ratings['movieId'] = movie_encoder.fit_transform(ratings.movieId)

user_encoder.classes_
movie_encoder.classes_

from scipy.sparse import csr_matrix
np.random.seed(42)
def create_matrix(data, user_col, item_col, rating_col):
    """
    creates the sparse user-item interaction matrix

    Parameters
    ----------
    data : DataFrame
        implicit rating data

    user_col : str
        user column name

    item_col : str
        item column name

    ratings_col : str
        implicit rating column name
    """
    
    data[[user_col, item_col]] = data[[user_col, item_col]].astype('category')
    
    rows = data[user_col].cat.codes
    cols = data[item_col].cat.codes
    rating = data[rating_col]
    user_item_matrix = csr_matrix((rating, (rows, cols)))
    return user_item_matrix

user_item_matrix = create_matrix(ratings, 'userId', 'movieId', 'rating')

user_item_matrix.shape

ratings['test'] = ratings.groupby(['userId'])['timestamp'].rank(method='first', ascending=False)

train_ratings = ratings.query('test != 1').drop(columns = ['test', 'timestamp'])
test_ratings = ratings.query('test == 1').drop(columns = ['test', 'timestamp'])

train_ratings.shape, test_ratings.shape

# bla = train_ratings.drop_duplicates(subset = ['userId', 'movieId'], keep = 'first')

# %%time
# unique_movies = set(train_ratings.movieId)
# def create_negative_movies(df, userid = 'userId', movieid = 'movieId',neg_examples = 4):
#     unique_movies = set(df[movieid])
    
#     movies = []
#     uids = df[userid].unique()
#     for u in uids:
#         movies.extend(np.random.choice(list(unique_movies - set(df[movieid][df[userid] == u])), size = neg_examples))
        
#     return uids, movies

# %%time 
# neg_examples = 4
# users, negative_movies = create_negative_movies(train_ratings)
# negative_movies_df = pd.DataFrame(dict(userId = np.repeat(users, [neg_examples]*len(users)),
#                 movieId = negative_movies,
#                 ratings = np.zeros(len(negative_movies)))
#                 )

# negative_movies_df.to_csv('negative_movies.csv', index = False)

print('Training Dimensions: ', train_ratings.userId.nunique(), train_ratings.movieId.nunique())
print('Test Dimensions: ', test_ratings.userId.nunique(), test_ratings.movieId.nunique())

print('Movies in Train: ', train_ratings.sum())
print('Movies in Test: ', test_ratings.sum())

train_users = train_ratings.userId.unique().tolist()
test_users = test_ratings.userId.unique().tolist()

print(len(train_users))
print(len(test_users))

def create_negative_df(user_ids, user_item, neg_examples = 4, test = False):
    
    movies_id = np.arange(user_item.shape[1])
    negative_movies = []
    examples = []
    for i in range(len(user_ids)):

        interacted = user_item[i].nonzero()[1]
        x = ~np.isin(movies_id, interacted)
        x = np.argwhere(x).squeeze(1)
        
        if test:
            size = neg_examples
        else:
            size = len(interacted)*neg_examples
        
        x = np.random.choice(x, size = size)
        negative_movies.extend(x)
        examples.append(size)
        
    negative_movies_df = pd.DataFrame(dict(userId = np.repeat(user_ids, examples),
                        movieId = negative_movies,
                        rating = np.zeros(len(negative_movies)))
                        )
    return negative_movies_df

%%time
train_negative_movies_df = create_negative_df(train_users, user_item_matrix, neg_examples = 4)
train_negative_movies_df.shape

%%time
test_negative_movies_df = create_negative_df(test_users, user_item_matrix, neg_examples = 99, test = True)
test_negative_movies_df.shape

full_training_df = train_ratings.append(train_negative_movies_df)
full_test_df = test_ratings.append(test_negative_movies_df)

full_training_df.shape, full_test_df.shape

full_training_df.info(memory_usage='deep'), full_test_df.info(memory_usage='deep')

full_training_df.columns

import torch
import torch.nn as nn
import pytorch_lightning as pl
from pytorch_lightning.callbacks import ModelCheckpoint

pl.seed_everything(42, workers=True)

from torch.utils.data import Dataset, DataLoader
from multiprocessing import cpu_count

class MovieData(Dataset):
    def __init__(self, users, movies, ratings):
        self.users = users
        self.movies = movies
        self.ratings = ratings
        
    def __len__(self):
        return len(self.ratings)
        
    def __getitem__(self, idx):
    
        users = self.users.iloc[idx]
        movies = self.movies.iloc[idx]
        ratings = self.ratings.iloc[idx]

        return dict(
            users = torch.tensor(users, dtype=torch.long),
            movies = torch.tensor(movies, dtype=torch.long),
            ratings = torch.tensor(ratings, dtype=torch.float)
        )

class MovieDataModule(pl.LightningDataModule):
    def __init__(self, train_df, test_df, batch_size = 512):
        super().__init__()
        
        self.train_df = train_df 
        self.test_df = test_df 
        self.batch_size = batch_size
        
    def setup(self, stage=None):
        
        self.train_data = MovieData(self.train_df.userId, self.train_df.movieId, self.train_df.rating)
        self.test_data = MovieData(self.test_df.userId, self.test_df.movieId, self.test_df.rating)
    
    def train_dataloader(self):
        return DataLoader(self.train_data, batch_size=self.batch_size, shuffle=True, pin_memory=True, num_workers = 10)
    
    def test_dataloader(self):
        return DataLoader(self.test_data, batch_size=self.batch_size, shuffle=False, pin_memory=True, num_workers = 10)
    

class NCF(nn.Module):
    def __init__(self, dim_users, dim_movies, n_out = 1):
        super().__init__()
        
        self.user_embedding = nn.Embedding(dim_users, 8)
        self.movie_embedding = nn.Embedding(dim_movies, 8)
        
        self.encoder = nn.Sequential(
                            nn.Linear(16,64),
                            nn.ReLU(inplace=True),
                            nn.Linear(64,32),
                            nn.ReLU(inplace=True),
                            nn.Linear(32,n_out)
                        )
        
    def forward(self, users, movies):
        user_emb = self.user_embedding(users)
        movie_emb = self.movie_embedding(movies)
        
        x = torch.cat((user_emb, movie_emb), dim = 1)
        x = self.encoder(x)
        return x
    
    
class RecSys(pl.LightningModule):
    def __init__(self, model):
        super().__init__()
        self.model = model
        self.criterion = nn.BCEWithLogitsLoss()
        
    def forward(self,users, movies):
        x = self.model(users, movies)
        return x
        
    def training_step(self, batch, batch_idx):
        users, movies, ratings = batch['users'], batch['movies'], batch['ratings']
        preds = self(users, movies)
        # print('preds:',  preds.shape)
        # print('ratings: ', ratings.shape)
        loss = self.criterion(preds, ratings.view(-1,1))
        self.log('train_loss', loss,  prog_bar = True, logger = True)
        return {'loss': loss}
    
    def configure_optimizers(self):
        return torch.optim.Adam(self.model.parameters(), lr = 1e-3)

full_test_df.userId.astype('int64').max(), full_test_df.movieId.astype('int64').max(), full_test_df.shape

dim_users = full_training_df.userId.astype('int64').max() + 1
dim_movies = full_training_df.movieId.astype('int64').max() + 1
print(dim_users, dim_movies)

model = NCF(dim_users, dim_movies)

dm = MovieDataModule(full_training_df, full_test_df, batch_size=512)
dm.setup()

train_batch = next(iter(dm.train_dataloader()))

value = np.random.randint(1,32)
train_batch['users'][value], train_batch['movies'][value], train_batch['ratings'][value]

train_batch['users'].shape, train_batch['movies'].shape, train_batch['ratings'].shape

recommender = RecSys(model)

recommender(train_batch['users'], train_batch['movies']).shape

mc = ModelCheckpoint(
    dirpath = 'checkpoints',
    #filename = 'best-checkpoint',
    save_last = True,
    save_top_k = 1,
    verbose = True,
    monitor = 'train_loss', 
    mode = 'min'
    )

mc.CHECKPOINT_NAME_LAST = 'best-checkpoint-latest'

trainer = pl.Trainer(max_epochs=5,
                    accelerator="gpu",
                    devices=1, 
                    callbacks=[mc], 
                    progress_bar_refresh_rate=30, 
                    # fast_dev_run=True,
                    #overfit_batches=1
                    )
trainer.fit(recommender, dm)

# from torchmetrics.functional import retrieval_hit_rate

# preds = torch.tensor([[0.9, 0.3, 0.9,0.4],[0.9, 0.3, 0.9,0.4]])
# target = torch.tensor([[False, False, True,False],[False, True, False,False]])
# retrieval_hit_rate(preds, target, k=2)

@torch.inference_mode()
def predict(model, dm):
    model.eval()
    preds = []
    for item in dm.test_dataloader():
        
        pred = torch.sigmoid(model(item['users'], item['movies']))
        preds.extend(pred.cpu().detach().numpy())
        
    return preds

predictions= np.array(predict(recommender, dm))
predictions.shape

full_test_df['preds'] = predictions
full_test_df

recomendations = full_test_df.sort_values(by = ['userId','preds'], ascending=[True, False]).groupby('userId').head(10)

# Hit Ratio @ 10

recomendations.rating.sum()/recomendations.userId.nunique()

def back_to_normal(df, user_encoder, movie_encoder, movies_mapping):
    
    idx_movies = df.movieId.tolist()
    idx_users = df.userId.tolist()
    return pd.DataFrame(dict(userId = user_encoder.classes_[idx_users],
                    movieId = pd.Series(movie_encoder.classes_[idx_movies]).map(movies_mapping),
                    rating = df.rating.tolist()))

visto= back_to_normal(train_ratings, user_encoder, movie_encoder, movies_mapping)
visto.shape

recomendar = back_to_normal(recomendations, user_encoder, movie_encoder, movies_mapping)
recomendar.shape

user = 4
print(visto.query('userId == @user')['movieId'])
recomendar.query('userId == @user')

user = 6265
print(visto.query('userId == @user')['movieId'])
recomendar.query('userId == @user')

user = 21962
print(visto.query('userId == @user')['movieId'])
recomendar.query('userId == @user')

user = 17568
print(visto.query('userId == @user')['movieId'])
recomendar.query('userId == @user')

user = 63
print(visto.query('userId == @user')['movieId'])
recomendar.query('userId == @user')

user = 162532
print(visto.query('userId == @user')['movieId'])
recomendar.query('userId == @user')