import streamlit as st
import fitz
import pandas as pd
import numpy as np
import re
import torch
from tqdm.auto import tqdm
from sentence_transformers import SentenceTransformer, util
import textwrap

# Set device
device = "cuda" if torch.cuda.is_available() else "cpu"

# Load PDF and process
@st.cache_data(show_spinner=True)
def process_pdf(pdf_path):
    def text_formatter(text):
        return text.replace("\n", " ").strip()

    doc = fitz.open(pdf_path)
    pages_and_texts = []
    for page_number, page in enumerate(doc):
        text = text_formatter(page.get_text())
        pages_and_texts.append({
            "page_number": page_number - 15,
            "text": text,
            "page_char_count": len(text),
            "page_word_count": len(text.split(" ")),
            "page_sentence_count_raw": len(text.split(". ")),
            "page_token_count": len(text)/4
        })
    return pages_and_texts

@st.cache_data(show_spinner=True)
def chunk_sentences(pages_and_texts, chunk_size=10):
    import spacy
    nlp = spacy.blank("en")
    nlp.add_pipe("sentencizer")

    def split_list(lst, size):
        return [lst[i:i+size] for i in range(0, len(lst), size)]

    for item in pages_and_texts:
        item["sentences"] = [str(s) for s in nlp(item["text"]).sents]
        item["sentence_chunks"] = split_list(item["sentences"], chunk_size)

    pages_and_chunks = []
    for item in pages_and_texts:
        for chunk in item["sentence_chunks"]:
            joined = "".join(chunk).replace("  ", " ").strip()
            joined = re.sub(r'\.([A-Z])', r'. \1', joined)
            pages_and_chunks.append({
                "page_number": item["page_number"],
                "sentence_chunk": joined,
                "chunk_char_count": len(joined),
                "chunk_word_count": len(joined.split(" ")),
                "chunk_token_count": len(joined)/4
            })
    return pages_and_chunks

@st.cache_data(show_spinner=True)
def embed_chunks(pages_and_chunks):
    model = SentenceTransformer("all-mpnet-base-v2", device=device)
    filtered = [c for c in pages_and_chunks if c["chunk_token_count"] > 30]
    texts = [item["sentence_chunk"] for item in filtered]
    embeddings = model.encode(texts, convert_to_tensor=True)
    for i, emb in enumerate(embeddings):
        filtered[i]["embedding"] = emb.cpu().numpy()
    return filtered, embeddings, model

@st.cache_data(show_spinner=False)
def semantic_search(query, embeddings, model, chunks, k=5):
    query_emb = model.encode(query, convert_to_tensor=True).to(device)
    scores = util.dot_score(query_emb, embeddings)[0]
    top_k = torch.topk(scores, k)
    results = []
    for score, idx in zip(top_k[0], top_k[1]):
        results.append({
            "score": float(score),
            "text": chunks[idx]["sentence_chunk"],
            "page": chunks[idx]["page_number"]
        })
    return results

st.title("Semantic Search App for Ian Goodfellow's Deep Learning Book")

uploaded_file = st.file_uploader("Upload the PDF", type="pdf")

if uploaded_file:
    with st.spinner("Processing PDF..."):
        with open("uploaded.pdf", "wb") as f:
            f.write(uploaded_file.read())
        texts = process_pdf("uploaded.pdf")
        chunks = chunk_sentences(texts)
        chunk_data, chunk_embeddings, emb_model = embed_chunks(chunks)

    query = st.text_input("Enter your query:")
    if query:
        with st.spinner("Searching..."):
            results = semantic_search(query, chunk_embeddings, emb_model, chunk_data)
        st.subheader("Top Results")
        for res in results:
            st.markdown(f"**Score:** {res['score']:.4f}")
            st.markdown(f"**Page:** {res['page']}")
            st.markdown(f"> {res['text']}")
            st.markdown("---")