import pandas as pd
import numpy as np
from zipfile import ZipFile
import tensorflow as tf
from tensorflow import keras
from pathlib import Path
import matplotlib.pyplot as plt
import gradio as gr
from huggingface_hub import from_pretrained_keras

# Download the actual data from http://files.grouplens.org/datasets/movielens/ml-latest-small.zip"
# Use the ratings.csv file
movielens_data_file_url = (
    "http://files.grouplens.org/datasets/movielens/ml-latest-small.zip"
)
movielens_zipped_file = keras.utils.get_file(
    "ml-latest-small.zip", movielens_data_file_url, extract=False
)
keras_datasets_path = Path(movielens_zipped_file).parents[0]
movielens_dir = keras_datasets_path / "ml-latest-small"

# Only extract the data the first time the script is run.
if not movielens_dir.exists():
    with ZipFile(movielens_zipped_file, "r") as zip:
        # Extract files
        print("Extracting all the files now...")
        zip.extractall(path=keras_datasets_path)
        print("Done!")

ratings_file = movielens_dir / "ratings.csv"
df = pd.read_csv(ratings_file)

# Make all the encodings
user_ids = df["userId"].unique().tolist()
user2user_encoded = {x: i for i, x in enumerate(user_ids)}
userencoded2user = {i: x for i, x in enumerate(user_ids)}
movie_ids = df["movieId"].unique().tolist()
movie2movie_encoded = {x: i for i, x in enumerate(movie_ids)}
movie_encoded2movie = {i: x for i, x in enumerate(movie_ids)}
df["user"] = df["userId"].map(user2user_encoded)
df["movie"] = df["movieId"].map(movie2movie_encoded)

num_users = len(user2user_encoded)
num_movies = len(movie_encoded2movie)
df["rating"] = df["rating"].values.astype(np.float32)
# min and max ratings will be used to normalize the ratings later
min_rating = min(df["rating"])
max_rating = max(df["rating"])

# Load model
model = from_pretrained_keras('mindwrapped/collaborative-filtering-movielens')
movie_df = pd.read_csv(movielens_dir / "movies.csv")


def update_user(id):
  return get_top_rated_from_user(id), get_recommendations(id)
  

def get_top_rated_from_user(id):
  decoded_id = userencoded2user.get(id)
  movies_watched_by_user = df[df.userId == decoded_id]

  # Get the top rated movies by this user
  top_movies_user = (
      movies_watched_by_user.sort_values(by="rating", ascending=False)
      .head(5)
      .movieId.values
  )
  movie_df_rows = movie_df[movie_df["movieId"].isin(top_movies_user)]
  movie_df_rows = movie_df_rows.drop('movieId', axis=1)
  return movie_df_rows


def random_user():
  return update_user(np.random.randint(0, num_users))


def get_recommendations(id):
  decoded_id = userencoded2user.get(id)
  movies_watched_by_user = df[df.userId == decoded_id]

  # Get the top 10 recommended movies for this user
  movies_not_watched = movie_df[
      ~movie_df["movieId"].isin(movies_watched_by_user.movieId.values)
  ]["movieId"]
  movies_not_watched = list(
      set(movies_not_watched).intersection(set(movie2movie_encoded.keys()))
  )
  movies_not_watched = [[movie2movie_encoded.get(x)] for x in movies_not_watched]

  # Encode user
  user_encoder = id

  # Create data [[user_id, movie_id],...]
  user_movie_array = np.hstack(
      ([[user_encoder]] * len(movies_not_watched), movies_not_watched)
  )

  # Predict ratings for movies not watched
  ratings = model.predict(user_movie_array).flatten()

  # Get indices of top ten movies
  top_ratings_indices = ratings.argsort()[-10:][::-1]

  # Decode each movie
  recommended_movie_ids = [
      movie_encoded2movie.get(movies_not_watched[x][0]) for x in top_ratings_indices
  ]
  recommended_movies = movie_df[movie_df["movieId"].isin(recommended_movie_ids)]
  recommended_movies = recommended_movies.drop('movieId', axis=1)

  return recommended_movies

demo = gr.Blocks()

with demo:
  with gr.Box():
    gr.Markdown(
    """
    ## Input
    #### Select a user to get recommendations for.
    """)

    inp1 = gr.Slider(0, num_users, value=0, label='User')
    # btn1 = gr.Button('Random User')

    # top_rated_from_user = get_top_rated_from_user(0)

    gr.Markdown(
    """
    #### Movies with the Highest Ratings from this user
    """)
    df1 = gr.DataFrame(interactive=False)

  with gr.Box():
    gr.Markdown('## Output')
    # recommendations = get_recommendations(0)
    gr.Markdown(
    """
    #### Top 10 movie recommendations
    """)
    df2 = gr.DataFrame(interactive=False)

  gr.HTML("""
  <p style='text-align: center; color: #C7C7C7'>
      <a href='https://keras.io/examples/structured_data/collaborative_filtering_movielens/' target='_blank' style='text-decoration: underline'>Keras Example by Siddhartha Banerjee</a>
      <br>
      Space by Scott Krstyen (mindwrapped)
      </p>
  """)
  
  
  inp1.change(fn=update_user,
              inputs=inp1,
              outputs=[df1, df2])
  

demo.launch()