import streamlit as st
from ctransformers import AutoModelForCausalLM
from fastapi import FastAPI
from pydantic import BaseModel
import numpy as np
import PyPDF2
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import time

# Load the model
@st.cache_resource
def load_model():
    return AutoModelForCausalLM.from_pretrained(
        "llama-2-7b-chat.Q5_K_M.gguf",
        model_type='llama',
        max_new_tokens=1096,
        threads=3,
    )

llm = load_model()

# Pydantic object
class Validation(BaseModel):
    prompt: str
    context: str

# FastAPI
app = FastAPI()

@app.post("/llm_on_cpu")
async def stream(item: Validation):
    system_prompt = "You are an intelligent assistant. Answer the question based solely on the provided context. If the answer cannot be found in the context, say 'I don't have enough information to answer this question based on the provided context.'"
    B_INST, E_INST = "[INST]", "[/INST]"
    B_SYS, E_SYS = "<<SYS>>\n", "\n<</SYS>>\n\n"
    
    full_prompt = f"{B_INST}{B_SYS}{system_prompt}{E_SYS}Context: {item.context}\n\nQuestion: {item.prompt}\n\nAnswer:{E_INST}"
    
    return llm(full_prompt, stream=True)

# Helper functions
@st.cache_resource
def get_vectorizer():
    return TfidfVectorizer()

def encode_texts(texts, vectorizer):
    return vectorizer.fit_transform(texts)

def retrieve_documents(query, documents, top_k=3):
    vectorizer = get_vectorizer()
    all_texts = documents + [query]
    embeddings = encode_texts(all_texts, vectorizer)
    doc_embeddings = embeddings[:-1]
    query_embedding = embeddings[-1]
    similarities = cosine_similarity(query_embedding, doc_embeddings).flatten()
    top_indices = np.argsort(similarities)[-top_k:][::-1]
    return [documents[i] for i in top_indices]

def generate_llama_response(prompt, context):
    system_prompt = "You are an intelligent assistant. Answer the question based solely on the provided context. If the answer cannot be found in the context, say 'I don't have enough information to answer this question based on the provided context.'"
    B_INST, E_INST = "[INST]", "[/INST]"
    B_SYS, E_SYS = "<<SYS>>\n", "\n<</SYS>>\n\n"
    
    full_prompt = f"{B_INST}{B_SYS}{system_prompt}{E_SYS}Context: {context}\n\nQuestion: {prompt}\n\nAnswer:{E_INST}"
    
    return llm(full_prompt, stream=True)

def stream_response(prompt, context, placeholder):
    response = ""
    for token in generate_llama_response(prompt, context):
        response += token
        if len(response.split()) >= 50:  # Minimum 50 tokens
            placeholder.markdown(response + "▌")
        if len(response.split()) >= 100:  # Maximum 100 tokens
            break
    placeholder.markdown(response)

def summarize_pdf(documents, placeholder):
    context = " ".join(documents[:3])
    prompt = "Provide a concise summary of the main points in the following text:"
    stream_response(prompt, context, placeholder)

def create_questions(documents, placeholder):
    context = " ".join(documents[:3])
    prompt = "Generate 5 important questions based on the following text. These questions should cover key concepts and encourage critical thinking:"
    stream_response(prompt, context, placeholder)

def create_notes(documents, placeholder):
    context = " ".join(documents[:3])
    prompt = "Create concise and informative notes based on the following text. Focus on key concepts, definitions, and important points:"
    stream_response(prompt, context, placeholder)

def extract_text_from_pdf(file):
    pdf_reader = PyPDF2.PdfReader(file)
    text = ""
    for page in pdf_reader.pages:
        text += page.extract_text()
    return text

# Streamlit UI
def main():
    st.title("LLaMA-2 PDF Analysis and Chat Interface")

    uploaded_file = st.file_uploader("Choose a PDF file", type="pdf")

    if uploaded_file is not None:
        text = extract_text_from_pdf(uploaded_file)
        chunk_size = 1000
        documents = [text[i:i+chunk_size] for i in range(0, len(text), chunk_size)]

        st.success("PDF uploaded and processed successfully!")

        # Create columns for buttons
        col1, col2, col3 = st.columns(3)

        with col1:
            if st.button("Summarize PDF"):
                with st.spinner("Generating summary..."):
                    placeholder = st.empty()
                    summarize_pdf(documents, placeholder)

        with col2:
            if st.button("Generate Questions"):
                with st.spinner("Generating questions..."):
                    placeholder = st.empty()
                    create_questions(documents, placeholder)

        with col3:
            if st.button("Create Notes"):
                with st.spinner("Creating notes..."):
                    placeholder = st.empty()
                    create_notes(documents, placeholder)

        # Chat interface
        st.subheader("Chat with the PDF")
        user_question = st.text_input("Ask a question about the PDF:")
        if user_question:
            relevant_docs = retrieve_documents(user_question, documents)
            context = " ".join(relevant_docs)
            with st.spinner("Generating answer..."):
                placeholder = st.empty()
                stream_response(user_question, context, placeholder)

    else:
        st.info("Please upload a PDF file to begin.")

if __name__ == "__main__":
    main()