import torch
from transformers import BartForConditionalGeneration, BartTokenizer
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import nltk
from nltk.tokenize import sent_tokenize

nltk.download('punkt')

class HuggingFaceQAAssistant:
    def __init__(self, model_name='facebook/bart-large-cnn'):
        self.knowledge_base = {}
        self.vectorizer = TfidfVectorizer()
        self.tfidf_matrix = None
        self.tokenizer = BartTokenizer.from_pretrained(model_name)
        self.model = BartForConditionalGeneration.from_pretrained(model_name)
        
    def load_file(self, file_path):
        try:
            with open(file_path, 'r', encoding='utf-8') as file:
                content = file.read()
                file_name = file_path.split('/')[-1]
                self.knowledge_base[file_name] = content
                print(f"Successfully loaded {file_name}")
        except FileNotFoundError:
            print(f"Error: File {file_path} not found.")
        except Exception as e:
            print(f"Error loading file: {str(e)}")
    
    def build_tfidf_matrix(self):
        documents = list(self.knowledge_base.values())
        self.tfidf_matrix = self.vectorizer.fit_transform(documents)
    
    def get_most_relevant_document(self, query):
        if self.tfidf_matrix is None:
            self.build_tfidf_matrix()
        
        query_vector = self.vectorizer.transform([query])
        similarities = cosine_similarity(query_vector, self.tfidf_matrix).flatten()
        most_similar_idx = similarities.argmax()
        return list(self.knowledge_base.values())[most_similar_idx]
    
    def get_relevant_context(self, document, query, max_length=1024):
        sentences = sent_tokenize(document)
        sentence_vectors = self.vectorizer.transform(sentences)
        query_vector = self.vectorizer.transform([query])
        
        similarities = cosine_similarity(query_vector, sentence_vectors).flatten()
        sorted_idx = similarities.argsort()[::-1]
        
        context = ""
        for idx in sorted_idx:
            if len(context) + len(sentences[idx]) <= max_length:
                context += sentences[idx] + " "
            else:
                break
        
        return context.strip()
    
    def answer_query(self, query):
        relevant_doc = self.get_most_relevant_document(query)
        context = self.get_relevant_context(relevant_doc, query)
        
        input_text = f"context: {context}\n\nquestion: {query}\n\nanswer:"
        input_ids = self.tokenizer.encode(input_text, return_tensors="pt", max_length=1024, truncation=True)
        
        with torch.no_grad():
            output = self.model.generate(input_ids, max_length=150, num_beams=4, early_stopping=True)
        
        answer = self.tokenizer.decode(output[0], skip_special_tokens=True)
        return answer
    
    def list_loaded_files(self):
        if self.knowledge_base:
            return "Loaded files:\n" + "\n".join(self.knowledge_base.keys())
        else:
            return "No files have been loaded yet."

# Ejemplo de uso
assistant = HuggingFaceQAAssistant()

# Bucle principal
while True:
    print("\nHugging Face Q&A Assistant")
    print("1. Load a file")
    print("2. Ask a question")
    print("3. List loaded files")
    print("4. Exit")
    
    choice = input("Enter your choice (1-4): ")
    
    if choice == '1':
        file_path = input("Enter the path to the text file: ")
        assistant.load_file(file_path)
    elif choice == '2':
        query = input("Enter your question: ")
        answer = assistant.answer_query(query)
        print("Answer:", answer)
    elif choice == '3':
        print(assistant.list_loaded_files())
    elif choice == '4':
        print("Goodbye!")
        break
    else:
        print("Invalid choice. Please try again.")