# 2. Import necessary libraries
from transformers import AutoModelForCausalLM, AutoTokenizer
import gradio as gr
import torch
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
import os

# 3. Load the base model and tokenizer from Hugging Face
model_name = "EleutherAI/gpt-neo-1.3B"  # You can replace this with another LLM available for free
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# 4. Function to tokenize, vectorize, and find similar parts of a document
def find_relevant_parts(document_text, prompt, top_n=5):
    vectorizer = TfidfVectorizer(stop_words='english')
    doc_parts = document_text.split('.')  # Split document into parts
    vectors = vectorizer.fit_transform(doc_parts)
    
    query_vec = vectorizer.transform([prompt])
    similarities = (vectors * query_vec.T).toarray().flatten()
    top_indices = similarities.argsort()[-top_n:][::-1]  # Top `top_n` similar parts
    relevant_parts = ' '.join([doc_parts[i] for i in top_indices])
    return relevant_parts

# 5. Function for generating a response from the model
def generate_response(document_path, prompt):
    with open(document_path, 'r') as file:
        document_text = file.read()
    
    relevant_parts = find_relevant_parts(document_text, prompt)
    complete_input = prompt + ' ' + relevant_parts
    inputs = tokenizer(complete_input, return_tensors="pt")
    outputs = model.generate(**inputs, max_length=150, num_return_sequences=1)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

# 6. Gradio Interface setup for user authentication and document querying
def user_authentication(username, password):
    # Mock simple authentication logic (Replace with secure user authentication if needed)
    valid_users = {"user1": "pass1", "user2": "pass2"}  # Replace with more secure methods
    if username in valid_users and valid_users[username] == password:
        return True
    return False

# 7. Gradio App Structure
with gr.Blocks() as demo:
    # Page 1: Login page
    with gr.Tab("Login"):
        username = gr.Textbox(label="Username")
        password = gr.Textbox(label="Password", type="password")
        login_button = gr.Button("Login")
        login_output = gr.Textbox(label="Login Status", interactive=False)
        
        def login_check(user, pwd):
            if user_authentication(user, pwd):
                return "Login successful!", gr.update(visible=False), gr.update(visible=True)
            else:
                return "Invalid credentials", gr.update(visible=True), gr.update(visible=False)
        
        login_button.click(login_check, inputs=[username, password], outputs=[login_output])

    # Page 2: Document selection
    with gr.Tab("Select Document", visible=False) as doc_tab:
        doc_path = gr.File(label="Upload your document")
        select_button = gr.Button("Select Document")
        doc_status = gr.Textbox(label="Document Status", interactive=False)

        def document_selected(file):
            if file:
                return "Document selected successfully!", gr.update(visible=False), gr.update(visible=True)
            return "No document selected."
        
        select_button.click(document_selected, inputs=[doc_path], outputs=[doc_status])

    # Page 3: Query input and response
    with gr.Tab("Query Document", visible=False) as query_tab:
        prompt = gr.Textbox(label="Enter your prompt")
        submit_button = gr.Button("Submit Query")
        response_output = gr.Textbox(label="Response", interactive=False)

        def process_query(file, user_prompt):
            if file and user_prompt:
                return generate_response(file.name, user_prompt)
            return "Please provide a valid document and prompt."
        
        submit_button.click(process_query, inputs=[doc_path, prompt], outputs=[response_output])

# Launch the Gradio app
demo.launch()