import streamlit as st
from transformers import AutoModelForSequenceClassification
import pandas as pd
import numpy as np
import torch
import pickle
import wikipedia

from preprocess import Preprocess
from utility import Utility

st.title("Movie Genre Predictor")

text = st.text_input('Enter plot of the movie')
st.caption("Either enter Wiki URL or the Cast info of the movie. Cast will be fetched from the Wiki page if cast is not provided")
wiki_url = st.text_input("Enter Wiki URL of the movie (Needed for fetching the cast information)")
cast_input = st.text_input("Enter Wiki IDs of the cast (Should be separated by comma)")

model = AutoModelForSequenceClassification.from_pretrained("./checkpoint-49092")
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model.to(device)

lr_model = pickle.load(open("models/cast_plot_lr","rb"))
cast_mlb = pickle.load(open("models/cast_mlb","rb"))
#column_names = pickle.load(open("models/column_names","rb"))
#top_actors = list(column_names)[11:]
top_actors = list(cast_mlb.classes_)

meta_model = pickle.load(open("models/meta_model","rb"))

utility = Utility()
preprocess = Preprocess()
out = []

if st.button("Predict"):
    cast = []
    if len(wiki_url)!=0 and len(cast_input)==0:
            html_page = wikipedia.page(title=wiki_url.split("/")[-1].replace("_"," "), auto_suggest=False).html()
            cast_wiki = html_page.split(" title=\"Edit section: Cast\">edit</a>")[-1]
            anchor_tags = cast_wiki.split("<a href=")[1:6]
            top5_cast_links = [val.split("\"")[1] for val in anchor_tags]
            for actor in top5_cast_links:
                  try:
                    cast.append(wikipedia.page(title=actor.split("/")[-1].replace("_"," ")).pageid)
                  except:
                    pass
    else:
            if ", " in cast_input:
                cast = cast_input.split(", ")
            else:
                cast = cast_input.split(",")
    
    cast_str = ""
    for actor in cast:
            cast_str += actor + ", "
    st.write("Wiki Ids of Top 5 Cast:",cast_str)
    st.write("Genre: ")
    
    clean_plot = preprocess.apply(text)

    # Use Meta Model approach when cast information is available otherwise use DistilBERT model
    if len(cast)!=0:
            # Base Model 1: DistilBERT
            id2label, label2id, tokenizer, tokenized_plot = utility.tokenize(clean_plot, ["Action","Drama", "Romance", "Comedy", "Thriller","Crime","Horror"])
            input_ids = [np.asarray(tokenized_plot['input_ids'])]
            attention_mask = [np.asarray(tokenized_plot['attention_mask'])]
        
            y_pred = model(torch.IntTensor(input_ids), torch.IntTensor(attention_mask))
            pred = torch.FloatTensor(y_pred['logits'][0])
            
            sigmoid = torch.nn.Sigmoid()
            distilbert_pred = sigmoid(pred.squeeze().cpu())
        
            # Base model 2: LR One Vs All
            cast_features = []
            for actor in cast:
                if actor in top_actors:
                    cast_features.append(str(actor))
            lr_model_pred = lr_model.predict_proba(cast_mlb.transform([cast_features]))
        
            # Concatenating Outputs of base models
            r1 = distilbert_pred[3]
            r2 = distilbert_pred[5]
            r3 = distilbert_pred[1]
            r4 = distilbert_pred[6]
            r5 = distilbert_pred[2]
            r6 = distilbert_pred[4]
            distilbert_pred[1] = r1
            distilbert_pred[2] = r2
            distilbert_pred[3] = r3
            distilbert_pred[4] = r4
            distilbert_pred[5] = r5
            distilbert_pred[6] = r6
            pred1 = distilbert_pred
            pred2 = lr_model_pred
            distilbert_pred = pred1.detach().numpy()
            lr_model_pred = np.array(pred2)[0]
            concat_features = np.concatenate((lr_model_pred,distilbert_pred))
        
            # Meta model 3: LR One Vs All
            probs = meta_model.predict_proba([concat_features])
            
            # Preparing Output
            id2label = {0: "Action",1: "Comedy",2: "Crime",3: "Drama",4: "Horror",5: "Romance",6: "Thriller"}
            i = 0
            for prob in probs[0]:
                out.append([id2label[i], prob])
                i += 1
    else:
         id2label, label2id, tokenizer, tokenized_plot = utility.tokenize(clean_plot, ["Action","Drama", "Romance", "Comedy", "Thriller","Crime","Horror"])
         input_ids = [np.asarray(tokenized_plot['input_ids'])]
         attention_mask = [np.asarray(tokenized_plot['attention_mask'])]
         
         y_pred = model(torch.IntTensor(input_ids), torch.IntTensor(attention_mask))
         pred = torch.FloatTensor(y_pred['logits'][0])
         
         sigmoid = torch.nn.Sigmoid()
         probs = sigmoid(pred.squeeze().cpu())
         
         i = 0
         for prob in probs:
                 out.append([id2label[i], np.asscalar(prob)])
                 i += 1   
    
    st.write(out)