import gradio as gr
from fastai.collab import *
from fastai.tabular.all import *
from fastai.learner import load_learner
import pandas as pd
import numpy as np

dados = pd.read_csv('valid.csv')
ids = dados['user'].unique()
ids_list = list(map(str, ids.tolist()))

ratings = pd.read_csv('ratings.csv')

class CollabNN(Module):
    def __init__(self, user_sz, item_sz, y_range=(0,5.5), n_act=100):
        self.user_factors = Embedding(*user_sz)
        self.item_factors = Embedding(*item_sz)
        self.layers = nn.Sequential(
            nn.Linear(user_sz[1]+item_sz[1], n_act),
            nn.ReLU(),
            nn.Linear(n_act, 1))
        self.y_range = y_range
        
    def forward(self, x):
        embs = self.user_factors(x[:,0]),self.item_factors(x[:,1])
        x = self.layers(torch.cat(embs, dim=1))
        return sigmoid_range(x, *self.y_range)

class DotProduct(Module):
    def __init__(self, n_users, n_books, n_factors, y_range=(0, 5.5)):
        self.user_factors = Embedding(n_users, n_factors)
        self.user_bias = Embedding(n_users, 1)
        self.books_factors = Embedding(n_books, n_factors)
        self.books_bias = Embedding(n_books, 1)
        self.y_range = y_range
    
    def forward(self, x):
        users = self.user_factors(x[:,0])
        books = self.books_factors(x[:,1])
        res = (users * books).sum(dim=1, keepdim=True)
        res += self.user_bias(x[:,0]) + self.books_bias(x[:,1])
        return sigmoid_range(res, *self.y_range)

learn = load_learner('model.pkl', cpu=True)

def top5(user):
    user = int(user)
    items = pd.Series(learn.dls.classes['title']).unique()
    
    clas_items = ratings.loc[(ratings['user'] == user) & (ratings['rating'] > 0), 'title']
    no_clas_items = np.setdiff1d(items, clas_items)

    df = pd.DataFrame({'user': [user]*len(no_clas_items), 'title': no_clas_items})
    
    preds,_ = learn.get_preds(dl=learn.dls.test_dl(df))
    
    df['prediction'] = preds.numpy()
    
    top_5 = df.nlargest(5, 'prediction')
    
    return '\n'.join(top_5['title'].tolist())

iface = gr.Interface(
    fn=top5, 
    inputs=gr.Dropdown(choices=ids_list), 
    outputs="text", 
    title="Books Recommendation",
    description="This model is responsible for a recommendation system involving books and their ratings.",
)

iface.launch(share=True)