import gradio as gr
import joblib
import pandas as pd
import heapq

# Load saved artifacts
svd = joblib.load('svd_model.pkl')
movies = joblib.load('movies.pkl')
ratings = joblib.load('ratings.pkl')

def recommend_movies(user_id, N=10):
    anti_testset = [(u, i, r) for (u, i, r) in ratings[['userId', 'movieId', 'rating']].itertuples(index=False) if u != user_id]
    user_anti_testset = [t for t in anti_testset if t[0] == user_id][:1000]  # Limit for speed
    predictions = svd.test(user_anti_testset)
    top_n = heapq.nlargest(N, predictions, key=lambda x: x.est)
    top_movie_ids = [pred.iid for pred in top_n]
    top_titles = movies[movies['movieId'].isin(top_movie_ids)]['title'].values
    return "\n".join(top_titles)

demo = gr.Interface(
    fn=recommend_movies,
    inputs=[gr.Number(label="User ID", value=1), gr.Number(label="N", value=10)],
    outputs="text",
    title="Movie Recommendation System",
    description="Enter a User ID and N to get top movie recommendations."
)
demo.launch()