import os
import numpy as np
import fitz  # PyMuPDF pour extraction PDF
import faiss
import pickle
import matplotlib.pyplot as plt
from concurrent.futures import ThreadPoolExecutor
from openai import OpenAI
from sklearn.manifold import TSNE
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from dotenv import load_dotenv
import seaborn as sns

# Charger les variables d'environnement
load_dotenv()
OPENAI_API_KEY = os.getenv('OPENAI_API_KEY_static')

# 📌 Initialisation du client OpenAI
client = OpenAI(api_key=OPENAI_API_KEY)
model_embedding = "text-embedding-ada-002"
model_chat = "gpt-4-turbo"
temperature = 0.1  # Réduction de la température pour privilégier la RAG

# 📌 Paramètres de segmentation
chunk_size = 256  # Réduction du chunk size pour un meilleur contrôle du contexte
chunk_overlap = 15

# 📌 Définition des chemins de stockage
index_path = "faiss_index_openai.bin"
chunks_path = "chunked_docs_openai.pkl"
metadata_path = "metadata_openai.pkl"
embeddings_path = "embeddings_openai.npy"

# 📌 Vérification et chargement des données
if os.path.exists(index_path) and os.path.exists(chunks_path) and os.path.exists(metadata_path) and os.path.exists(embeddings_path):
    print("🔄 Chargement des données existantes...")
    index = faiss.read_index(index_path)
    with open(chunks_path, "rb") as f:
        chunked_docs = pickle.load(f)
    with open(metadata_path, "rb") as f:
        metadata_list = pickle.load(f)
    embeddings = np.load(embeddings_path)
    print("✅ Index, chunks, embeddings et métadonnées chargés avec succès !")
else:
    print("⚡ Création et stockage d'un nouvel index FAISS...")
    
    # 📌 Extraction des documents et métadonnées
    def extract_and_chunk_pdfs(pdf_folder):
        documents = SimpleDirectoryReader(pdf_folder, recursive=True).load_data()
        chunked_docs, metadata_list = [], []
        for doc in documents:
            doc_text = doc.text
            file_name = doc.metadata.get("file_name", "Inconnu")
            title = doc.metadata.get("title") or os.path.splitext(file_name)[0]
            doc_metadata = {"source": file_name, "title": title}
            for i in range(0, len(doc_text), chunk_size):
                chunk = doc_text[i:i + chunk_size]
                chunked_docs.append({"text": chunk, "metadata": doc_metadata})
                metadata_list.append(doc_metadata)
        return chunked_docs, metadata_list
    
    pdf_folder = 'C:/Users/MIPO10053340/OneDrive - Groupe Avril/Bureau/Salon_Agriculture_2024/Micka_API_Call/Docs_pdf/'
    chunked_docs, metadata_list = extract_and_chunk_pdfs(pdf_folder)
    
    # 📌 Génération des embeddings en parallèle
    def get_embeddings_in_batches(text_chunks, batch_size=5):
        embeddings = []
        
        def process_batch(batch):
            response = client.embeddings.create(input=[chunk["text"] for chunk in batch], model=model_embedding)
            return [data.embedding for data in response.data]
        
        with ThreadPoolExecutor(max_workers=5) as executor:
            future_batches = [executor.submit(process_batch, text_chunks[i:i+batch_size]) for i in range(0, len(text_chunks), batch_size)]
            for future in future_batches:
                embeddings.extend(future.result())
        
        return np.array(embeddings).astype('float32')
    
    embeddings = get_embeddings_in_batches(chunked_docs)
    
    # 📌 Création et stockage de l'index FAISS
    dimension = embeddings.shape[1]
    index = faiss.IndexFlatL2(dimension)
    index.add(embeddings)
    faiss.write_index(index, index_path)
    
    # 📌 Sauvegarde des données
    with open(chunks_path, "wb") as f:
        pickle.dump(chunked_docs, f)
    with open(metadata_path, "wb") as f:
        pickle.dump(metadata_list, f)
    np.save(embeddings_path, embeddings)  
    print("✅ Index, chunks, embeddings et métadonnées sauvegardés !")

# 📌 Récupération des chunks les plus pertinents
def retrieve_relevant_chunks(question, k=5):
    question_embedding_response = client.embeddings.create(
        input=[question],
        model=model_embedding
    )
    question_embedding = np.array(question_embedding_response.data[0].embedding).astype('float32').reshape(1, -1)
    distances, indices = index.search(question_embedding, k)
    if len(indices[0]) == 0:
        print("⚠️ Aucun chunk pertinent trouvé.")
        return []
    return [chunked_docs[i] for i in indices[0]]

# 📌 Génération de réponse avec OpenAI
def generate_response(context, question, sources):
    chunk_references = [f"[{i+1}]" for i in range(len(sources))]
    chunk_texts = "\n\n".join([f"{chunk_references[i]} (Source: {src['metadata']['source']}) :\n{src['text']}" for i, src in enumerate(sources)])
    messages = [
        {"role": "system", "content": f"Voici les informations extraites des documents :\n{chunk_texts}\n\nUtilise ces informations pour répondre."},
        {"role": "user", "content": question}
    ]
    response = client.chat.completions.create(
        model=model_chat,
        messages=messages,
        temperature=temperature
    )
    return response.choices[0].message.content + " " + "".join(chunk_references), chunk_texts

# 📌 Exécuter une requête utilisateur
user_question = "Quels sont les besoins en protéines des poulets de chair en phase de croissance ?"
relevant_chunks = retrieve_relevant_chunks(user_question)
context = "\n".join([chunk["text"] for chunk in relevant_chunks])
answer, citations = generate_response(context, user_question, relevant_chunks)

# 💾 Sauvegarde des résultats
with open("openai_response.txt", "w", encoding="utf-8") as f:
    f.write(f"Question : {user_question}\n")
    f.write(f"Réponse :\n{answer}\n")
    f.write(f"{citations}\n")

print("\n✅ Réponse enregistrée dans 'openai_response.txt'")

# 📊 Visualisation des embeddings avec t-SNE : 
    
# Nos documents ont un vocabulaire et un contenu très similaires, donc les embeddings générés par Mistral sont proches les uns des autres.
# t-SNE réduit la dimensionnalité, mais si les embeddings de base sont très proches, la distinction entre eux est moins visible en 2D.
# Un clustering plus clair apparaîtrait si nos documents couvraient des thématiques très variées (ex: alimentation, maladies, croissance de différents animaux).

# Génération de la réduction de dimension avec t-SNE
tsne = TSNE(n_components=2, perplexity=min(30, max(2, embeddings.shape[0] - 1)), random_state=42)
embeddings_2d = tsne.fit_transform(embeddings)

# Récupération des étiquettes des sources
source_labels = [chunk['metadata']['source'] for chunk in chunked_docs]

# Création du graphique
plt.figure(figsize=(10, 8))
sns.scatterplot(x=embeddings_2d[:, 0], y=embeddings_2d[:, 1], hue=source_labels, palette='tab10', alpha=0.7)
plt.title('Visualisation des embeddings avec t-SNE')
plt.xlabel('Dimension 1')
plt.ylabel('Dimension 2')
# Limiter la légende aux 10 premières sources
handles, labels = plt.gca().get_legend_handles_labels()
plt.legend(handles[:10], labels[:10], title="Fichiers sources (Top 10)", loc="upper right", bbox_to_anchor=(1.2, 1))


# Sauvegarde du graphique en PNG
output_path = "C:/Users/MIPO10053340/OneDrive - Groupe Avril/Bureau/Salon_Agriculture_2024/Micka_API_Call/embeddings_visualization.png"
plt.savefig(output_path, dpi=300, bbox_inches='tight')