import gradio as gr
import pandas as pd

from fastai.tabular.all import *
from fastai.layers import Module, Embedding, sigmoid_range

import torch
import torch.nn.functional as F

class DotProductBias(Module):
    def __init__(self, n_users, n_games, n_factors, y_range=(0.1, 11754.0)):
        self.user_factors = Embedding(n_users, n_factors)
        self.user_bias = Embedding(n_users, 1)
        self.game_factors = Embedding(n_games, n_factors)
        self.game_bias = Embedding(n_games, 1)
        self.y_range = y_range

    def forward(self, x):
        users = self.user_factors(x[:, 0])
        games = self.game_factors(x[:, 1])
        res = (users * games).sum(dim=1, keepdim=True)
        res += self.user_bias(x[:, 0]) + self.game_bias(x[:, 1])
        return torch.sigmoid(res) * (self.y_range[1] - self.y_range[0]) + self.y_range[0]


df = pd.read_csv('original.csv')

path = Path("./model.pkl")
learn = load_learner(path, 'model.pkl')


examples = list(df['user-id'].unique())[:20]

all_games_title = list(df['game-title'].unique())


def predict(user_id):

    user_id = int(user_id)

    games_alredy_played = list(df[(df['user-id'] == user_id)]['game-title'].unique())

    games_not_played = set(all_games_title) - set(games_alredy_played)
    size_games_not_played = len(games_not_played)

    data = {
            'user-id': [user_id]* size_games_not_played,
            'game-title': list(games_not_played),
            'time-played': [0]* size_games_not_played
        }

    new_df = pd.DataFrame(data)
    new_dl = learn.dls.test_dl(new_df)
    predictions = learn.get_preds(dl=new_dl)
    predicted_time_played = predictions[0].squeeze().tolist()
    new_df['time-played'] = predicted_time_played
    recomendations = new_df.sort_values(by='time-played', ascending=False).head(10)['game-title'].tolist()
    st = [r + '\n' for r in recomendations]

    ans =  ''.join(st)
    return ans


gr.Interface(
    fn=predict,
    inputs="text",
    outputs="text",
    examples=examples
).launch(share=False)