import streamlit as st
import pandas as pd
import numpy as np
import tensorflow as tf
import joblib
import os

# Define file paths (assuming all files are in the same folder as this script)
BASE_DIR = os.path.dirname(__file__)
KERAS_MODEL_PATH = os.path.join(BASE_DIR, "recommender_model.keras")
MOVIES_PATH = os.path.join(BASE_DIR, "movies.csv")
ENCODINGS_PATH = os.path.join(BASE_DIR, "encodings.pkl")

@st.cache_resource
def load_model():
    if not os.path.exists(KERAS_MODEL_PATH):
        st.error(f"❌ Model file not found at: {KERAS_MODEL_PATH}")
        st.stop()
    try:
        return tf.keras.models.load_model(KERAS_MODEL_PATH)
    except Exception as e:
        st.error(f"❌ Failed to load model:\n\n{e}")
        st.stop()

@st.cache_data
def load_assets():
    if not os.path.exists(MOVIES_PATH):
        st.error("❌ movies.csv not found.")
        st.stop()
    if not os.path.exists(ENCODINGS_PATH):
        st.error("❌ encodings.pkl not found.")
        st.stop()
    try:
        df_movies = pd.read_csv(MOVIES_PATH)
        user_map, movie_map = joblib.load(ENCODINGS_PATH)
        return df_movies, user_map, movie_map
    except Exception as e:
        st.error(f"❌ Failed to load assets:\n\n{e}")
        st.stop()

# Load model and assets
model = load_model()
movies_df, user2idx, movie2idx = load_assets()
reverse_movie_map = {v: k for k, v in movie2idx.items()}

# UI
st.title("🎬 TensorFlow Movie Recommender")
st.write("Select some movies you've liked to get personalized recommendations:")

# Movie title selection
movie_titles = movies_df.set_index("movieId")["title"].to_dict()
movie_choices = [movie_titles[mid] for mid in movie2idx if mid in movie_titles]
selected_titles = st.multiselect("🎞️ Liked movies", sorted(movie_choices))

# User ratings dict
user_ratings = {}
for title in selected_titles:
    movie_id = next((k for k, v in movie_titles.items() if v == title), None)
    if movie_id:
        user_ratings[movie_id] = 5.0

# Generate recommendations
if st.button("🎯 Get Recommendations"):
    if not user_ratings:
        st.warning("Please select at least one movie.")
    else:
        liked_indices = [movie2idx[m] for m in user_ratings if m in movie2idx]
        if not liked_indices:
            st.error("⚠️ No valid movie encodings found.")
            st.stop()

        try:
            # Calculate average embedding and similarity scores
            avg_embedding = tf.reduce_mean(model.layers[2](tf.constant(liked_indices)), axis=0, keepdims=True)
            all_movie_indices = tf.range(len(movie2idx))
            movie_embeddings = model.layers[3](all_movie_indices)
            scores = tf.reduce_sum(avg_embedding * movie_embeddings, axis=1).numpy()
            top_indices = np.argsort(scores)[::-1]

            # Top 10 recommendations excluding already liked
            recommended = []
            for idx in top_indices:
                mid = reverse_movie_map.get(idx)
                if mid not in user_ratings and mid in movie_titles:
                    recommended.append((movie_titles[mid], scores[idx]))
                if len(recommended) >= 10:
                    break

            st.subheader("🍿 Top 10 Recommendations")
            for title, score in recommended:
                st.write(f"**{title}** — Score: `{score:.3f}`")
        except Exception as e:
            st.error(f"❌ Error generating recommendations:\n\n{e}")