import os
import streamlit as st

from datasets import load_dataset
import re
import time
import requests
from bs4 import BeautifulSoup
import pandas as pd
import torch.nn.functional as F
from torch import Tensor
from transformers import AutoTokenizer, AutoModel
from keybert import KeyBERT
from keyphrase_vectorizers import KeyphraseCountVectorizer
kw_model=KeyBERT(model='AI-Growth-Lab/PatentSBERTa')



tokenizer = AutoTokenizer.from_pretrained('intfloat/e5-base')
model = AutoModel.from_pretrained('intfloat/e5-base')


# Function to fetch categories, title, and related text from a Wikipedia page
def fetch_wikipedia_data(article_title):
    url = f"https://en.wikipedia.org/wiki/{article_title.replace(' ', '_')}"

    response = requests.get(url)

    if response.status_code == 200:
        soup = BeautifulSoup(response.text, 'html.parser')

        # Find the categories section at the bottom of the page
        categories_section = soup.find("div", {"class": "mw-normal-catlinks"})

        if categories_section:
            # Extract individual categories
            categories = [cat.text for cat in categories_section.find("ul").find_all("li")]

            # Extract the title
            title = article_title

            return {"title": title, "categories": categories}

    return None

def get_wiki_category_aprch_1(pat_text):
    print(pat_text)
    keywords=kw_model.extract_keywords(pat_text,keyphrase_ngram_range=(1, 3),top_n=15,vectorizer=KeyphraseCountVectorizer())
    titles=[]
    for i in range(len(keywords)):
        title=keywords[i][0]
        titles.append(title)
    data = []
    for i in titles:
        results = fetch_wikipedia_data(i)
        data.append(results)
    cats=[]
    for i in range(len(data)):
        if data[i] is not None:
            cat=data[i]['categories']
            cats.append(cat)
    result=[j for i in cats for j in i]
    res = [i for n, i in enumerate(result) if i not in result[:n]]
    return titles,res


#     return res

def average_pool(last_hidden_states: Tensor,
                attention_mask: Tensor) -> Tensor:
    last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
    return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
            
def get_wiki_category(pat_text):
    # print(pat_text)
    keywords=kw_model.extract_keywords(pat_text,keyphrase_ngram_range=(1, 3),top_n=3,vectorizer=KeyphraseCountVectorizer())
    titles=[]
    for i in range(len(keywords)):
        title=keywords[i][0]
        titles.append(title)

    batch_dict = tokenizer(titles, padding=True, truncation=True, return_tensors='pt')
    outputs = model(**batch_dict)
    embeddings = average_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
    embeddings = F.normalize(embeddings, p=2, dim=1)
    values = embeddings.tolist()

    catgories_list = []
    for value in values:
        try:
            response = index.query(vector=value,top_k=3,include_metadata=True)
        except:
            pinecone.init(api_key='b33ddf5d-5b1a-4d0e-9a3f-572008563791',environment='gcp-starter')
            index = pinecone.Index("wiki-index")
            response = index.query(vector=value,top_k=5,include_metadata=True)

        catgories = response['matches'][0]['metadata']['categories']
        catgories_list.append(catgories.split(','))
    flatList = [element for innerList in catgories_list for element in innerList]
    new_list = [item.replace("'", '') for item in flatList]
    a_list = [s.strip() for s in new_list]
    test_list = list(set(a_list))
    # result=[j for i in flatList for j in i]
    # res = [i for n, i in enumerate(result) if i not in result[:n]]
    return test_list

def main():
    st.title('wikipedia-titles-category-generator')
    
    text = st.text_area("Enter a text paragraph:")

    if st.button('Get Wiki categories'):
        st.write("Predicting Wiki Categories for text:",text[:200])
        start_time = time.time()
        titles,wiki_categories=get_wiki_category_aprch_1(text)
        end_time = time.time()
        st.write({f"Wiki_titles for Text":titles})
        st.write({f"Wiki_categories for Text":wiki_categories})
            
if __name__ == "__main__":
    main()