import os
import json
import time
import re
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import accuracy_score, f1_score, confusion_matrix
from llama_cpp import Llama

# Set style for premium charts
sns.set_theme(style="darkgrid")
plt.rcParams.update({
    'grid.color': '#2c2c2c',
    'axes.facecolor': '#121212',
    'figure.facecolor': '#0e0e0e',
    'text.color': '#e0e0e0',
    'axes.labelcolor': '#e0e0e0',
    'xtick.color': '#b0b0b0',
    'ytick.color': '#b0b0b0',
    'font.size': 11
})

print("=== [Process 2] Evaluating Gemma 3 LLM & Plotting Comparisons ===", flush=True)

# 1. Load intermediate results from Keras baseline
if not os.path.exists("data/keras_results.csv"):
    raise FileNotFoundError("data/keras_results.csv not found! Please run evaluate/evaluate_keras.py first.")

df_sample = pd.read_csv("data/keras_results.csv")
print(f"Loaded {len(df_sample)} pre-evaluated tweets from Keras baseline.", flush=True)

# 2. Loading Gemma 3 GGUF Model
print("Loading local Gemma 3 GGUF model via Llama.from_pretrained...", flush=True)
llm = Llama.from_pretrained(
    repo_id="JusteLeo/emotion-text-classifier-LLM",
    filename="EmotionTextClassifierLLM.gguf",
    n_ctx=512,
    verbose=False
)
print("Gemma 3 GGUF loaded successfully.", flush=True)

# Emotion Mapping Definitions
positive_emotions = {
    'joy', 'love', 'surprise', 'pride', 'admiration', 'gratitude', 'hope', 
    'optimism', 'amusement', 'desire', 'caring', 'relief', 'excitement', 
    'approval', 'caring', 'curiosity'
}

negative_emotions = {
    'sadness', 'anger', 'fear', 'disgust', 'shame', 'guilt', 'disappointment', 
    'annoyance', 'frustration', 'grief', 'nervousness', 'embarrassment', 
    'remorse', 'disapproval', 'confusion', 'boredom'
}

def clean_and_parse_json(text):
    """Robustly cleans and parses JSON from the LLM output."""
    cleaned = text.strip()
    cleaned = re.sub(r"^```(?:json)?", "", cleaned, flags=re.IGNORECASE)
    cleaned = re.sub(r"```$", "", cleaned).strip()
    
    try:
        data = json.loads(cleaned)
        return data
    except Exception:
        emotions = []
        explanation = "Error parsing explanation"
        
        emotion_match = re.search(r'"emotions"\s*:\s*\[(.*?)\]', cleaned, re.DOTALL)
        if emotion_match:
            emotions_str = emotion_match.group(1)
            emotions = [e.strip(' "\'') for e in emotions_str.split(',')]
            
        explanation_match = re.search(r'"explanation"\s*:\s*"(.*?)"', cleaned, re.DOTALL)
        if explanation_match:
            explanation = explanation_match.group(1)
            
        return {"emotions": emotions, "explanation": explanation}

llm_emotions = []
llm_explanations = []
llm_predictions = []
llm_times = []

# Run LLM inference loop
print("Starting Gemma 3 predictions loop (500 tweets)...", flush=True)
for idx, text in enumerate(df_sample['text']):
    if (idx + 1) % 25 == 0 or idx == 0:
        print(f"Processing tweet {idx+1}/{len(df_sample)}...", flush=True)
        
    start_time = time.time()
    
    try:
        response = llm.create_chat_completion(
            messages=[{"role": "user", "content": text}],
            temperature=0.1,
            max_tokens=80
        )
        output_content = response["choices"][0]["message"]["content"]
        parsed = clean_and_parse_json(output_content)
        
        emotions = parsed.get("emotions", [])
        explanation = parsed.get("explanation", "No explanation.")
        
        # Map primary emotion to binary sentiment
        mapped_sentiment = 0 # Default negative
        if emotions:
            primary_emotion = emotions[0].lower().strip()
            if primary_emotion in positive_emotions:
                mapped_sentiment = 1
            elif primary_emotion in negative_emotions:
                mapped_sentiment = 0
            else:
                pos_count = sum(1 for e in emotions if e.lower().strip() in positive_emotions)
                neg_count = sum(1 for e in emotions if e.lower().strip() in negative_emotions)
                if pos_count > neg_count:
                    mapped_sentiment = 1
        else:
            emotions = ["Neutral"]
            explanation = "Unable to classify fine emotions."
            mapped_sentiment = 0
            
    except Exception as e:
        print(f"Error at index {idx}: {str(e)}", flush=True)
        emotions = ["Error"]
        explanation = f"Inference failed with error: {str(e)}"
        mapped_sentiment = 0
        
    end_time = time.time()
    
    llm_emotions.append(", ".join(emotions))
    llm_explanations.append(explanation)
    llm_predictions.append(mapped_sentiment)
    llm_times.append(end_time - start_time)

df_sample['llm_emotions'] = llm_emotions
df_sample['llm_explanation'] = llm_explanations
df_sample['llm_pred'] = llm_predictions
df_sample['llm_time'] = llm_times

# Calculate metrics
keras_accuracy = accuracy_score(df_sample['true_sentiment'], df_sample['keras_pred'])
keras_f1 = f1_score(df_sample['true_sentiment'], df_sample['keras_pred'])
keras_avg_time = df_sample['keras_time'].iloc[0]

llm_accuracy = accuracy_score(df_sample['true_sentiment'], df_sample['llm_pred'])
llm_f1 = f1_score(df_sample['true_sentiment'], df_sample['llm_pred'])
llm_total_time = sum(llm_times)
llm_avg_time = np.mean(llm_times)

print(f"\nBaseline Keras Accuracy: {keras_accuracy:.4f}", flush=True)
print(f"Baseline Keras F1-Score: {keras_f1:.4f}", flush=True)
print(f"Baseline Keras Avg Speed: {keras_avg_time*1000:.2f} ms / tweet", flush=True)

print(f"\nGemma 3 Accuracy: {llm_accuracy:.4f}", flush=True)
print(f"Gemma 3 F1-Score: {llm_f1:.4f}", flush=True)
print(f"Gemma 3 Avg Speed: {llm_avg_time:.4f} s / tweet", flush=True)

# Save final benchmark results
df_sample.to_csv("data/benchmark_results.csv", index=False)
print("\nFinal benchmark results saved to data/benchmark_results.csv", flush=True)

# 3. Generate Visualizations
print("\n--- Generating Premium Performance Visualizations ---", flush=True)
os.makedirs("assets", exist_ok=True)

# 3.1 Confusion Matrices
fig, axes = plt.subplots(1, 2, figsize=(14, 6))

# Keras Matrix
cm_keras = confusion_matrix(df_sample['true_sentiment'], df_sample['keras_pred'])
sns.heatmap(cm_keras, annot=True, fmt='d', cmap='Blues', cbar=False, ax=axes[0],
            xticklabels=['Négatif', 'Positif'], yticklabels=['Négatif', 'Positif'],
            annot_kws={"size": 14, "weight": "bold"})
axes[0].set_title("Matrice de Confusion : Baseline Keras", fontsize=14, weight='bold', pad=15)
axes[0].set_xlabel("Prédiction", fontsize=12)
axes[0].set_ylabel("Vrai Label", fontsize=12)

# Gemma Matrix
cm_llm = confusion_matrix(df_sample['true_sentiment'], df_sample['llm_pred'])
sns.heatmap(cm_llm, annot=True, fmt='d', cmap='Greens', cbar=False, ax=axes[1],
            xticklabels=['Négatif', 'Positif'], yticklabels=['Négatif', 'Positif'],
            annot_kws={"size": 14, "weight": "bold"})
axes[1].set_title("Matrice de Confusion : Challengeur Gemma 3 (LLM)", fontsize=14, weight='bold', pad=15)
axes[1].set_xlabel("Prédiction", fontsize=12)
axes[1].set_ylabel("Vrai Label", fontsize=12)

plt.tight_layout()
plt.savefig("assets/confusion_matrices.png", dpi=300, facecolor='#0e0e0e')
plt.close()

# 3.2 Accuracy Bar Chart
fig, ax = plt.subplots(figsize=(8, 6))
models = ['Keras (Baseline)', 'Gemma 3 (LLM)']
accuracies = [keras_accuracy * 100, llm_accuracy * 100]
speeds = [keras_avg_time, llm_avg_time]

color_acc = '#00df89'
bars = ax.bar(models, accuracies, width=0.5, color=color_acc, alpha=0.85, edgecolor='#00df89', linewidth=1.5)
ax.set_ylabel('Précision (Accuracy %)', color=color_acc, fontsize=12, weight='bold')
ax.set_ylim(0, 100)
ax.tick_params(axis='y', labelcolor=color_acc)
ax.set_xticklabels(models, fontsize=12, weight='bold')

for bar in bars:
    height = bar.get_height()
    ax.text(bar.get_x() + bar.get_width()/2., height + 2, f'{height:.1f}%', ha='center', va='bottom', color='#00df89', weight='bold', fontsize=12)

plt.title("Comparaison de la Justesse de Prédiction (Accuracy)", fontsize=14, weight='bold', pad=20, color='white')
plt.tight_layout()
plt.savefig("assets/accuracy_comparison.png", dpi=300, facecolor='#0e0e0e')
plt.close()

# 3.3 Speed Bar Chart
fig, ax = plt.subplots(figsize=(8, 6))
color_speed = '#ff5e62'
bars_time = ax.bar(models, speeds, width=0.5, color=color_speed, alpha=0.85, edgecolor='#ff5e62', linewidth=1.5)
ax.set_ylabel("Temps d'inférence moyen par tweet (s)", color=color_speed, fontsize=12, weight='bold')
ax.tick_params(axis='y', labelcolor=color_speed)
ax.set_xticklabels(models, fontsize=12, weight='bold')

for bar in bars_time:
    height = bar.get_height()
    time_text = f'{height*1000:.1f} ms' if height < 0.01 else f'{height:.3f} s'
    ax.text(bar.get_x() + bar.get_width()/2., height + 0.02 * max(speeds), time_text, ha='center', va='bottom', color='#ff5e62', weight='bold', fontsize=12)

plt.title("Comparaison de la Vitesse d'Inférence", fontsize=14, weight='bold', pad=20, color='white')
plt.tight_layout()
plt.savefig("assets/speed_comparison.png", dpi=300, facecolor='#0e0e0e')
plt.close()

# 4. Construct Qualitative Explainability Table
print("\n--- Constructing Qualitative Explainability Table ---", flush=True)
ex_happy = df_sample[(df_sample['true_sentiment'] == 1) & (df_sample['llm_pred'] == 1)].head(2)
ex_sad = df_sample[(df_sample['true_sentiment'] == 0) & (df_sample['llm_pred'] == 0)].head(2)
ex_disagree = df_sample[df_sample['keras_pred'] != df_sample['llm_pred']].head(1)

if len(ex_disagree) == 0:
    ex_disagree = df_sample.head(1)

qualitative_sample = pd.concat([ex_happy, ex_sad, ex_disagree]).head(5)

qual_markdown = """### Table 1: Analyse Qualitative de l'Explicabilité (Explainability) de Gemma 3

Ce tableau présente 5 exemples de tweets analysés lors de la preuve de concept, confrontant les prédictions binaires de la baseline Keras avec l'analyse d'émotion fine et l'explication formulée par le LLM Gemma 3.

| Tweet Original | Vrai Sentiment | Prédiction Keras (Probabilité) | Émotion Fine détectée (Gemma 3) | Explication Générée par Gemma 3 |
| :--- | :---: | :---: | :---: | :--- |
"""

for _, row in qualitative_sample.iterrows():
    true_lbl = "😊 Positif" if row['true_sentiment'] == 1 else "😢 Négatif"
    keras_lbl = f"Positif ({row['keras_prob']:.2f})" if row['keras_pred'] == 1 else f"Négatif ({row['keras_prob']:.2f})"
    
    tweet_txt = row['text'].replace('|', '\\|').replace('\n', ' ')
    explanation_txt = row['llm_explanation'].replace('|', '\\|').replace('\n', ' ')
    
    qual_markdown += f"| \"{tweet_txt}\" | {true_lbl} | {keras_lbl} | **{row['llm_emotions']}** | {explanation_txt} |\n"

with open("assets/qualitative_table.md", "w", encoding="utf-8") as f:
    f.write(qual_markdown)

print("Visualizations and qualitative table successfully generated in the assets directory!", flush=True)
print("=== [Process 2] Finished successfully ===", flush=True)