# app.py
import gradio as gr
from datasets import load_dataset
import random
import pandas as pd
import faiss
import numpy as np
from sentence_transformers import SentenceTransformer, InputExample, losses
from torch.utils.data import DataLoader

# ---------------------------
# 1️⃣ Load datasets
# ---------------------------
movies = load_dataset("AiresPucrs/movielens-movies", split="train")
ratings = load_dataset("AiresPucrs/movielens-user-ratings", split="train")

# Use full dataset
movies_list = [dict(m) for m in movies]
movie_titles = [m['title'] for m in movies_list]

# Combine title + genre for embeddings
movie_texts = [f"{m['title']} — {m['genres']}" for m in movies_list]

# Ratings DataFrame
df_ratings = pd.DataFrame(ratings)

# Safe movieId -> title mapping
movieId_to_title = {m['movieId']: m['title'] for m in movies_list}

# ---------------------------
# 2️⃣ Generate triplets for fine-tuning
# ---------------------------
triplets = []

for m in movies_list:
    genres = set(m['genres'].split("|"))
    
    # Positive: movies sharing at least 1 genre
    same_genre = [x for x in movies_list if x != m and len(set(x['genres'].split("|")).intersection(genres)) > 0]
    if not same_genre:
        continue
    pos = random.choice(same_genre)
    
    # Negative: movies with no genre overlap
    diff_genre = [x for x in movies_list if len(set(x['genres'].split("|")).intersection(genres)) == 0]
    if not diff_genre:
        continue
    neg = random.choice(diff_genre)
    
    triplets.append((
        f"{m['title']} — {m['genres']}",
        f"{pos['title']} — {pos['genres']}",
        f"{neg['title']} — {neg['genres']}"
    ))

# ---------------------------
# 3️⃣ Fine-tune embeddings (fast: 1 epoch)
# ---------------------------
model = SentenceTransformer('all-MiniLM-L6-v2')
train_examples = [InputExample(texts=[a,p,n]) for a,p,n in triplets]
train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=16)
train_loss = losses.TripletLoss(model=model)

# Fine-tune (1 epoch for testing)
model.fit(train_objectives=[(train_dataloader, train_loss)], epochs=1, warmup_steps=50)

# ---------------------------
# 4️⃣ Build FAISS index
# ---------------------------
movie_embeddings = model.encode(movie_texts, show_progress_bar=True)
movie_embeddings = np.array(movie_embeddings).astype("float32")
dim = movie_embeddings.shape[1]
index = faiss.IndexFlatL2(dim)
index.add(movie_embeddings)

# ---------------------------
# 5️⃣ Recommendation functions
# ---------------------------
def recommend_similar_movies(query_title, top_k=5):
    movie_row = [m for m in movies_list if m['title'] == query_title]
    if len(movie_row) == 0:
        return ["Movie not found"]
    
    query_text = f"{movie_row[0]['title']} — {movie_row[0]['genres']}"
    q_emb = model.encode([query_text]).astype("float32")
    
    # Search more neighbors to ensure enough recommendations
    D, I = index.search(q_emb, top_k * 2)
    
    recommendations = []
    for idx in I[0]:
        title = movie_titles[int(idx)]
        if title != query_title and title not in recommendations:
            recommendations.append(title)
        if len(recommendations) >= top_k:
            break
            
    return recommendations

def recommend_for_user(user_id, top_k=5):
    user_rated = df_ratings[df_ratings["userId"] == user_id]
    liked = user_rated[user_rated["rating"] >= 4]["movieId"].tolist()
    
    similar_users = df_ratings[df_ratings["movieId"].isin(liked)]["userId"].unique()
    recs = df_ratings[
        (df_ratings["userId"].isin(similar_users)) &
        (~df_ratings["movieId"].isin(liked)) &
        (df_ratings["rating"] >= 4)
    ]["movieId"]
    
    top_recs = recs.value_counts().head(top_k).index.tolist()
    recommended_titles = [movieId_to_title[m] for m in top_recs if m in movieId_to_title]
    return recommended_titles if recommended_titles else ["No recommendations available"]

# ---------------------------
# 6️⃣ Gradio interface
# ---------------------------
def gradio_recommend(input_type, value, top_k=5):
    if input_type == "Movie Title":
        return recommend_similar_movies(value, top_k)
    elif input_type == "User ID":
        try:
            user_id = int(value)
            return recommend_for_user(user_id, top_k)
        except:
            return ["Invalid User ID"]
    else:
        return ["Invalid input type"]

iface = gr.Interface(
    fn=gradio_recommend,
    inputs=[
        gr.Dropdown(["Movie Title", "User ID"], label="Input Type"),
        gr.Textbox(label="Enter movie title or user ID"),
        gr.Slider(1, 10, value=5, step=1, label="Number of recommendations")
    ],
    outputs=gr.JSON(label="Recommendations"),  # ✅ JSON ensures correct counts
    title="Movie Recommendation System",
    description="Content-based and user-personalized movie recommendations"
)

iface.launch()