import streamlit as st
from fastai.collab import *
import torch
from torch import nn
import pickle
import pandas as pd
from transformers import PegasusForConditionalGeneration, PegasusTokenizer
import sentencepiece
import string
import requests

@st.cache_resource
def load_stuff():
  # Load the data loader 
  dls= pd.read_pickle('dataloader.pkl')
  # Create an instance of the model
  learn = collab_learner(dls, use_nn=True,layers=[20,10],y_range=(0,10.5))
  # Load the saved state dictionary
  state_dict = torch.load('myModel.pth',map_location=torch.device('cpu'))
  # Assign the loaded state dictionary to the model's load_state_dict() method
  learn.model.load_state_dict(state_dict)   
  #load books dataframe
  books = pd.read_csv('./data/BX_Books.csv', sep=';',encoding='latin-1')
  #load tokenizer
  tokenizer = PegasusTokenizer.from_pretrained("pszemraj/pegasus-x-large-book-summary")
  #load model
  model = PegasusForConditionalGeneration.from_pretrained("pszemraj/pegasus-x-large-book-summary")
  return dls, learn, books, tokenizer, model

dls, learn, books, tokenizer, model = load_stuff()

#function to get recommendations
def get_3_recs(book):
  book_factors = learn.model.embeds[1].weight
  idx = dls.classes['title'].o2i[book]
  distances = nn.CosineSimilarity(dim=1)(book_factors, book_factors[idx][None])
  idxs = distances.argsort(descending=True)[1:4]
  recs = [dls.classes['title'][i] for i in idxs]
  return recs

#function to get descriptions from Google Books
def search_book_description(title):
  # Google Books API endpoint for book search
  url = "https://www.googleapis.com/books/v1/volumes"
   # Parameters for the book search
  params = {
    "q": title,
    "maxResults": 1
    }
  # Send GET request to Google Books API
  response = requests.get(url, params=params)
  # Check if the request was successful
  if response.status_code == 200:
    # Parse the JSON response to extract the book description
      data = response.json()

      if "items" in data and len(data["items"]) > 0:
        book_description = data["items"][0]["volumeInfo"].get("description", "No description available.")
        return book_description
      else:
        print("No book found with the given title.")
        return None
  else:
    # If the request failed, print the error message
    print("Error:", response.status_code, response.text)
    return None
  
#function to ensure summaries end with punctuation
def cut(sum):
    last_punc_idx = max(sum.rfind(p) for p in string.punctuation)
    output = sum[:last_punc_idx + 1]
    return output


#function to summarize
def summarize(des_list):
    if "No description available." in des_list:
      idx = des_list.index("No description available.")
      des = des_list.copy()
      des.pop(idx)
      rest = summarize(des)
      rest.insert(idx,'No description available.')
      return rest
    else: 
      # Tokenize all the descriptions in the list
      encoded_inputs = tokenizer(des_list, truncation=True, padding="longest", return_tensors="pt")

      # Generate summaries for all the inputs
      summaries = model.generate(**encoded_inputs, max_new_tokens=100)

      # Decode the summaries and process them
      outputs = tokenizer.batch_decode(summaries, skip_special_tokens=True)
      outputs = list(map(cut, outputs))
      return outputs

#function to get cover images  
def get_covers(recs):
  imgs = [books[books['Book-Title']==r]['Image-URL-L'].tolist()[0]for r in recs]
  return imgs

#streamlit app construction
st.title('Your digital librarian')
st.markdown("Hi there! I recommend you books based on one you love (which might not be in the same genre because that's boring) and give you my own synopsis of each book. Enjoy!")
options = books["Book-Title"].tolist()
input = st.selectbox('Select your favorite book', options)
if st.button("Get recommendations"):
   recs = get_3_recs(input)
   descriptions = list(map(search_book_description,recs))
   des_sums = summarize(descriptions)
   imgs = get_covers(recs)
   
   col1, col2, col3 = st.columns(3)
   col1.image(imgs[0])
   col1.markdown(f"**{recs[0]}**")
   col1.write(des_sums[0])

   col2.image(imgs[1])
   col2.markdown(f"**{recs[1]}**")
   col2.write(des_sums[1])

   col3.image(imgs[2])
   col3.markdown(f"**{recs[2]}**")
   col3.write(des_sums[2])