# ==============================================================================
#  AetherMap — ABLATION STUDY v3 (Rigoroso)
#  SQuAD v2 (EN) com perguntas sem resposta, intervalos de confiança,
#  análise de significância estatística e Pareto score vs latência
#  Para rodar no Google Colab contra a API do HF Space
# ==============================================================================

# %%
# !pip install requests pandas numpy matplotlib datasets scipy -q

# %%
import os
import json
import time
import random
import requests
import pandas as pd
import numpy as np
import tempfile
from typing import List, Dict, Any, Tuple
from datasets import load_dataset

# ==============================================================================
# CONFIGURAÇÃO
# ==============================================================================
OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY", "")
AETHERMAP_URL = "https://madras1-aethermap.hf.space"

EVALUATOR_MODEL = "nvidia/nemotron-3-nano-30b-a3b:free"

DELAY_BETWEEN_LLM_CALLS = 2.0

OPENROUTER_HEADERS = {
    "Authorization": f"Bearer {OPENROUTER_API_KEY}",
    "Content-Type": "application/json",
}

ABLATION_MODES = {
    "faiss_only":     {"label": "FAISS Only",          "faiss": True,  "bm25": False, "rrf": False, "reranker": False, "expansion": False},
    "bm25_only":      {"label": "BM25 Only",           "faiss": False, "bm25": True,  "rrf": False, "reranker": False, "expansion": False},
    "hybrid":         {"label": "Hybrid (RRF)",        "faiss": True,  "bm25": True,  "rrf": True,  "reranker": False, "expansion": False},
    "hybrid_rerank":  {"label": "+ Reranker",          "faiss": True,  "bm25": True,  "rrf": True,  "reranker": True,  "expansion": False},
    "full":           {"label": "Full Pipeline",       "faiss": True,  "bm25": True,  "rrf": True,  "reranker": True,  "expansion": True},
}

# Configurações do estudo
N_QUERIES = 50           # Mais queries = mais poder estatístico
N_UNANSWERABLE = 10      # Queries do SQuAD v2 SEM resposta
N_CONTEXTS = 300         # Contextos para indexar
N_BOOTSTRAP = 1000       # Iterações de bootstrap para IC
CONFIDENCE_LEVEL = 0.95  # Nível de confiança (95%)

print("✅ Configuração carregada")
print(f"🌐 API: {AETHERMAP_URL}")
print(f"🧑‍⚖️ Avaliador: {EVALUATOR_MODEL}")
print(f"📚 Dataset: SQuAD v1.1 + v2 (English)")
print(f"📏 N = {N_QUERIES} queries ({N_QUERIES - N_UNANSWERABLE} answerable + {N_UNANSWERABLE} unanswerable)")
print(f"📊 Bootstrap: {N_BOOTSTRAP} iterações, {CONFIDENCE_LEVEL*100:.0f}% CI")
if not OPENROUTER_API_KEY:
    raise RuntimeError("Defina OPENROUTER_API_KEY no ambiente antes de rodar o estudo de ablação.")

# %%
# ==============================================================================
# CARREGAR SQUAD v1 + v2
# ==============================================================================

def carregar_squad(
    n_contexts: int = 300,
    n_answerable: int = 40,
    n_unanswerable: int = 10,
    seed: int = 42,
) -> Tuple[List[Dict], List[Dict]]:
    """
    Carrega SQuAD v1.1 (answerable) + SQuAD v2 (unanswerable).
    
    Mistura questions com resposta e SEM resposta para testar se o RAG
    sabe dizer "não sei" quando o documento não contém a informação.
    """
    random.seed(seed)
    
    # --- SQuAD v1 para contextos + Q&A answerable ---
    print("📥 Baixando SQuAD v1.1 (answerable)...")
    ds_v1 = load_dataset("rajpurkar/squad", split="validation")
    print(f"   📊 SQuAD v1: {len(ds_v1)} exemplos")
    
    # Extrair contextos únicos
    context_set = {}
    for item in ds_v1:
        ctx = item["context"].strip()
        if ctx and len(ctx) > 50 and ctx not in context_set:
            context_set[ctx] = {
                "text": ctx,
                "title": item.get("title", ""),
            }
    
    all_contexts = list(context_set.values())
    print(f"   📄 {len(all_contexts)} contextos únicos")
    
    # Amostrar contextos
    sampled_contexts = random.sample(all_contexts, min(n_contexts, len(all_contexts)))
    sampled_texts = set(c["text"] for c in sampled_contexts)
    
    # Coletar Q&A answerable
    qa_answerable = []
    for item in ds_v1:
        ctx = item["context"].strip()
        if ctx in sampled_texts:
            answers = item.get("answers", {})
            answer_texts = answers.get("text", [])
            if answer_texts and answer_texts[0].strip():
                qa_answerable.append({
                    "pergunta": item["question"].strip(),
                    "resposta": answer_texts[0].strip(),
                    "has_answer": True,
                    "contexto_preview": ctx[:150],
                })
    
    # Diversificar por contexto
    random.shuffle(qa_answerable)
    seen = set()
    diverse_answerable = []
    for qa in qa_answerable:
        key = qa["contexto_preview"][:80]
        if key not in seen:
            diverse_answerable.append(qa)
            seen.add(key)
        if len(diverse_answerable) >= n_answerable:
            break
    
    if len(diverse_answerable) < n_answerable:
        remaining = [q for q in qa_answerable if q not in diverse_answerable]
        diverse_answerable.extend(remaining[:n_answerable - len(diverse_answerable)])
    
    print(f"   ✅ {len(diverse_answerable)} queries answerable selecionadas")
    
    # --- SQuAD v2 para unanswerable ---
    print("📥 Baixando SQuAD v2 (unanswerable)...")
    ds_v2 = load_dataset("rajpurkar/squad_v2", split="validation")
    print(f"   📊 SQuAD v2: {len(ds_v2)} exemplos")
    
    # Coletar queries SEM resposta
    qa_unanswerable = []
    for item in ds_v2:
        answers = item.get("answers", {})
        answer_texts = answers.get("text", [])
        # SQuAD v2: unanswerable = answers vazio
        if not answer_texts or not answer_texts[0].strip():
            qa_unanswerable.append({
                "pergunta": item["question"].strip(),
                "resposta": "",  # Sem resposta
                "has_answer": False,
                "contexto_preview": item["context"][:150],
            })
    
    random.shuffle(qa_unanswerable)
    selected_unanswerable = qa_unanswerable[:n_unanswerable]
    print(f"   ✅ {len(selected_unanswerable)} queries unanswerable selecionadas")
    
    # --- Combinar e embaralhar ---
    all_queries = diverse_answerable + selected_unanswerable
    random.shuffle(all_queries)
    
    print(f"\n   📋 Total: {len(all_queries)} queries ({len(diverse_answerable)} answerable + {len(selected_unanswerable)} unanswerable)")
    
    return sampled_contexts, all_queries


# %%
# ==============================================================================
# FUNÇÕES AUXILIARES
# ==============================================================================

def call_llm(prompt: str, model: str, max_tokens: int = 1000) -> str:
    """Chama LLM via OpenRouter usando requests direto."""
    time.sleep(DELAY_BETWEEN_LLM_CALLS)
    try:
        payload = {
            "model": model,
            "messages": [{"role": "user", "content": prompt}],
            "max_tokens": max_tokens,
            "temperature": 0.3,
            "reasoning": {"exclude": True},
        }
        resp = requests.post(
            "https://openrouter.ai/api/v1/chat/completions",
            headers=OPENROUTER_HEADERS,
            json=payload,
            timeout=90,
        )
        if resp.status_code != 200:
            print(f"   ⚠️ LLM ({model}): HTTP {resp.status_code}")
            return ""
        data = resp.json()
        choices = data.get("choices", [])
        if not choices:
            return ""
        content = choices[0].get("message", {}).get("content") or ""
        return content.strip()
    except Exception as e:
        print(f"   ⚠️ Erro LLM ({model}): {type(e).__name__}: {e}")
        time.sleep(5)
        return ""


def avaliar_resposta(pergunta: str, resposta_rag: str, resposta_esperada: str, has_answer: bool) -> Dict[str, float]:
    """
    Avaliação RIGOROSA:
    - Se has_answer=True: avaliar correcao, completude, relevancia
    - Se has_answer=False: avaliar se RAG disse "não sei"/"não encontrei"
      (penalizar alucinação — inventar resposta quando não deveria)
    """
    if has_answer:
        prompt = f"""Avalie a resposta do sistema RAG comparando com a referência humana.
Seja RIGOROSO — somente dê 1.0 se a informação estiver factualmente correta E completa.

PERGUNTA: {pergunta}
RESPOSTA DE REFERÊNCIA: {resposta_esperada[:300]}
RESPOSTA DO RAG: {resposta_rag[:500]}

Critérios (0.0 a 1.0):
- correcao: A resposta contém a mesma informação factual da referência? (0.0 = errado, 0.5 = parcial, 1.0 = correto)
- completude: Cobre TODOS os pontos? (0.0 = nada, 0.5 = metade, 1.0 = completo)
- relevancia: Responde DIRETAMENTE à pergunta sem enrolação? (0.0 = off-topic, 0.5 = tangencial, 1.0 = direto)

Retorne APENAS JSON: {{"correcao": X.X, "completude": X.X, "relevancia": X.X}}"""
    else:
        # Unanswerable — o RAG deveria dizer "não encontrei"
        prompt = f"""A pergunta abaixo NÃO TEM resposta nos documentos fornecidos.
Avalie se o sistema RAG CORRETAMENTE identificou que não pode responder.

PERGUNTA: {pergunta}
RESPOSTA DO RAG: {resposta_rag[:500]}

Critérios (0.0 a 1.0):
- correcao: O sistema disse que não encontrou/não pode responder? (1.0 = corretamente disse "não sei", 0.0 = inventou uma resposta falsa/alucinou)
- completude: Se recusou, foi claro e honesto? (1.0 = claro, 0.0 = confuso)
- relevancia: A resposta é apropriada para uma pergunta sem resposta? (1.0 = recusa educada, 0.0 = alucinação total)

Retorne APENAS JSON: {{"correcao": X.X, "completude": X.X, "relevancia": X.X}}"""

    response = call_llm(prompt, EVALUATOR_MODEL, max_tokens=500)

    try:
        if "{" in response:
            json_start = response.find("{")
            json_end = response.rfind("}")
            if json_end > json_start:
                result = json.loads(response[json_start:json_end + 1])
                return {
                    "correcao": max(0.0, min(1.0, float(result.get("correcao", 0.5)))),
                    "completude": max(0.0, min(1.0, float(result.get("completude", 0.5)))),
                    "relevancia": max(0.0, min(1.0, float(result.get("relevancia", 0.5)))),
                }
    except Exception as e:
        print(f"   ⚠️ Erro parsing avaliação: {e}")
    return {"correcao": 0.5, "completude": 0.5, "relevancia": 0.5}


# %%
# ==============================================================================
# ESTATÍSTICA
# ==============================================================================

def bootstrap_ci(scores: List[float], n_bootstrap: int = 1000, ci: float = 0.95) -> Tuple[float, float, float]:
    """
    Calcula intervalo de confiança via bootstrap.
    Returns: (mean, lower_bound, upper_bound)
    """
    if not scores:
        return 0.0, 0.0, 0.0
    
    arr = np.array(scores)
    mean = np.mean(arr)
    
    # Bootstrap
    boot_means = []
    for _ in range(n_bootstrap):
        sample = np.random.choice(arr, size=len(arr), replace=True)
        boot_means.append(np.mean(sample))
    
    boot_means = np.array(boot_means)
    alpha = 1 - ci
    lower = np.percentile(boot_means, 100 * alpha / 2)
    upper = np.percentile(boot_means, 100 * (1 - alpha / 2))
    
    return float(mean), float(lower), float(upper)


def compute_effect_size(scores_a: List[float], scores_b: List[float]) -> float:
    """Cohen's d — tamanho do efeito entre dois grupos."""
    a, b = np.array(scores_a), np.array(scores_b)
    pooled_std = np.sqrt((np.std(a, ddof=1)**2 + np.std(b, ddof=1)**2) / 2)
    if pooled_std == 0:
        return 0.0
    return float((np.mean(a) - np.mean(b)) / pooled_std)


# %%
# ==============================================================================
# UPLOAD E BUSCA
# ==============================================================================

def upload_contexts(contexts: List[Dict], text_column: str = "texto") -> str:
    """Faz upload dos contextos SQuAD para o AetherMap."""
    df = pd.DataFrame({
        "texto": [c["text"] for c in contexts],
        "titulo": [c["title"] for c in contexts],
    })
    with tempfile.NamedTemporaryFile(mode='w', suffix='.csv', delete=False, encoding='utf-8') as f:
        df.to_csv(f, index=False)
        temp_path = f.name
    
    print(f"📤 Upload de {len(contexts)} contextos...")
    with open(temp_path, "rb") as f:
        response = requests.post(
            f"{AETHERMAP_URL}/process/",
            files={"file": ("squad_contexts.csv", f)},
            data={"n_samples": len(contexts), "text_column": text_column, "fast_mode": "true"},
            timeout=300,
        )
    os.unlink(temp_path)
    
    if response.status_code != 200:
        raise Exception(f"Erro no upload: {response.text[:200]}")
    
    result = response.json()
    job_id = result["job_id"]
    meta = result["metadata"]
    print(f"✅ Job criado: {job_id[:8]}...")
    print(f"   📊 {meta['num_documents_processed']} docs | {meta['num_clusters_found']} clusters")
    return job_id


def search_with_mode(job_id: str, query: str, mode: str) -> Dict:
    """Executa busca com um modo de ablation específico."""
    start = time.time()
    try:
        response = requests.post(
            f"{AETHERMAP_URL}/search/",
            data={"query": query, "job_id": job_id, "ablation_mode": mode},
            timeout=300,
        )
        latency = time.time() - start
        if response.status_code != 200:
            return {"error": response.text[:100], "latency": latency}
        result = response.json()
        result["latency"] = latency
        return result
    except requests.exceptions.Timeout:
        latency = time.time() - start
        print(f"   ⏰ Timeout após {latency:.0f}s")
        return {"error": f"Timeout ({latency:.0f}s)", "latency": latency}
    except Exception as e:
        latency = time.time() - start
        return {"error": str(e)[:100], "latency": latency}


# %%
# ==============================================================================
# ABLATION STUDY PRINCIPAL
# ==============================================================================

def run_ablation_study():
    """Executa ablation study rigoroso."""
    print("=" * 70)
    print("🧪 AETHERMAP ABLATION STUDY v3 — RIGOROUS EVALUATION")
    print("=" * 70)
    
    n_answerable = N_QUERIES - N_UNANSWERABLE
    
    # 1. Carregar dados
    print("\n📚 FASE 1: Carregando SQuAD v1 + v2")
    contexts, queries = carregar_squad(N_CONTEXTS, n_answerable, N_UNANSWERABLE)
    
    # 2. Upload
    print(f"\n📦 FASE 2: Upload dos contextos")
    job_id = upload_contexts(contexts)
    
    # 3. Mostrar queries
    print(f"\n❓ FASE 3: {len(queries)} queries")
    for i, q in enumerate(queries):
        tag = "✓" if q["has_answer"] else "✗"
        print(f"   [{i+1}] [{tag}] Q: {q['pergunta'][:65]}...")
        if q["has_answer"]:
            print(f"            A: {q['resposta'][:60]}...")
        else:
            print(f"            A: (sem resposta nos docs)")
    
    # 4. Rodar modos
    print(f"\n🧪 FASE 4: Testando {len(ABLATION_MODES)} modos...")
    
    all_results = {}
    raw_scores = {}  # Guardar scores individuais pra estatística
    
    for mode_key, mode_info in ABLATION_MODES.items():
        mode_label = mode_info["label"]
        print(f"\n{'─' * 50}")
        print(f"🔬 Modo: {mode_label} ({mode_key})")
        components = [k.upper() for k, v in [("faiss", mode_info["faiss"]), ("bm25", mode_info["bm25"]),
                       ("rrf", mode_info["rrf"]), ("reranker", mode_info["reranker"]),
                       ("expansion", mode_info["expansion"])] if v]
        print(f"   Componentes: {' + '.join(components)}")
        
        scores = []
        latencies = []
        per_query_scores = []  # Score médio por query (pra bootstrap)
        answerable_scores = []
        unanswerable_scores = []
        
        for i, q in enumerate(queries):
            pergunta = q["pergunta"]
            esperada = q["resposta"]
            has_answer = q["has_answer"]
            
            result = search_with_mode(job_id, pergunta, mode_key)
            
            if "error" in result:
                print(f"   [{i+1}] ❌ Erro: {result['error'][:50]}")
                scores.append({"correcao": 0, "completude": 0, "relevancia": 0})
                latencies.append(result.get("latency", 0))
                per_query_scores.append(0.0)
                continue
            
            resposta = result.get("summary", "")
            latency = result.get("latency", 0)
            latencies.append(latency)
            
            aval = avaliar_resposta(pergunta, resposta, esperada, has_answer)
            scores.append(aval)
            
            avg = (aval["correcao"] + aval["completude"] + aval["relevancia"]) / 3
            per_query_scores.append(avg)
            
            if has_answer:
                answerable_scores.append(avg)
            else:
                unanswerable_scores.append(avg)
            
            tag = "✓" if has_answer else "✗"
            emoji = "✓" if avg >= 0.6 else "○"
            print(f"   [{i+1}] [{tag}] {emoji} C:{aval['correcao']:.2f} Com:{aval['completude']:.2f} R:{aval['relevancia']:.2f} | {latency:.1f}s")
            print(f"         Q: {pergunta[:75]}")
            if has_answer:
                print(f"         RAG: {resposta[:120]}...")
                if avg < 0.3:
                    print(f"         ⚠️ Esperado: {esperada[:80]}")
            else:
                rag_preview = resposta[:120] if resposta else "(vazio)"
                refused = avg >= 0.6
                print(f"         RAG: {rag_preview}...")
                print(f"         {'✅ Recusou corretamente' if refused else '❌ ALUCINOU (deveria recusar)'}")
            
            time.sleep(0.5)
        
        # --- Estatísticas ---
        avg_correcao = np.mean([s["correcao"] for s in scores])
        avg_completude = np.mean([s["completude"] for s in scores])
        avg_relevancia = np.mean([s["relevancia"] for s in scores])
        avg_score = (avg_correcao + avg_completude + avg_relevancia) / 3
        avg_latency = np.mean(latencies)
        std_latency = np.std(latencies)
        
        # Bootstrap CI
        mean_ci, lower_ci, upper_ci = bootstrap_ci(per_query_scores, N_BOOTSTRAP, CONFIDENCE_LEVEL)
        
        # Scores separados
        ans_mean = np.mean(answerable_scores) if answerable_scores else 0
        unans_mean = np.mean(unanswerable_scores) if unanswerable_scores else 0
        
        all_results[mode_key] = {
            "label": mode_label,
            "avg_score": avg_score,
            "correcao": avg_correcao,
            "completude": avg_completude,
            "relevancia": avg_relevancia,
            "avg_latency": avg_latency,
            "std_latency": std_latency,
            "ci_lower": lower_ci,
            "ci_upper": upper_ci,
            "answerable_score": ans_mean,
            "unanswerable_score": unans_mean,
            "n_queries": len(queries),
            "n_answerable": len(answerable_scores),
            "n_unanswerable": len(unanswerable_scores),
            "components": components,
        }
        raw_scores[mode_key] = per_query_scores
        
        ci_str = f"[{lower_ci:.3f}, {upper_ci:.3f}]"
        print(f"   📊 Score: {avg_score:.3f} {ci_str} | Ans:{ans_mean:.3f} Unans:{unans_mean:.3f} | Lat:{avg_latency:.1f}±{std_latency:.1f}s")
    
    # 5. Resultados
    print_results(all_results, raw_scores)
    plot_results(all_results)
    
    return all_results


# %%
# ==============================================================================
# RESULTADOS
# ==============================================================================

def print_results(results: Dict, raw_scores: Dict):
    """Imprime tabela comparativa com estatísticas."""
    print("\n" + "=" * 90)
    print("📊 RESULTADO DO ABLATION STUDY v3 — RIGOROUS")
    print("=" * 90)

    # Tabela principal
    print(f"\n{'Modo':<16} {'Score':>6} {'95% CI':>16} {'Correção':>8} {'Compl':>6} {'Relev':>6} {'Ans':>5} {'Unans':>5} {'Lat(s)':>8}")
    print("─" * 90)

    baseline_score = results.get("faiss_only", {}).get("avg_score", 0)
    best_score = max(r["avg_score"] for r in results.values())

    for mode_key in ABLATION_MODES.keys():
        if mode_key not in results:
            continue
        r = results[mode_key]
        ci_str = f"[{r['ci_lower']:.3f},{r['ci_upper']:.3f}]"
        trophy = " 🏆" if r["avg_score"] == best_score else ""
        lat_str = f"{r['avg_latency']:.1f}±{r['std_latency']:.1f}"

        print(f"{r['label']:<16} {r['avg_score']:>6.3f} {ci_str:>16} {r['correcao']:>8.3f} {r['completude']:>6.3f} {r['relevancia']:>6.3f} {r['answerable_score']:>5.3f} {r['unanswerable_score']:>5.3f} {lat_str:>8}{trophy}")

    # Effect sizes
    print(f"\n{'─' * 90}")
    print("📐 EFFECT SIZES (Cohen's d):")
    
    mode_pairs = [
        ("faiss_only", "hybrid", "FAISS → Hybrid"),
        ("hybrid", "hybrid_rerank", "Hybrid → +Reranker"),
        ("hybrid_rerank", "full", "+Reranker → Full"),
        ("faiss_only", "full", "FAISS → Full Pipeline"),
    ]
    
    for mode_a, mode_b, label in mode_pairs:
        if mode_a in raw_scores and mode_b in raw_scores:
            d = compute_effect_size(raw_scores[mode_b], raw_scores[mode_a])
            magnitude = "negligible" if abs(d) < 0.2 else "small" if abs(d) < 0.5 else "medium" if abs(d) < 0.8 else "large"
            print(f"   {label:<25} d = {d:+.3f} ({magnitude})")

    # Insights
    print(f"\n{'─' * 90}")
    print("💡 INSIGHTS:")
    
    if "faiss_only" in results and "bm25_only" in results:
        faiss = results["faiss_only"]["avg_score"]
        bm25 = results["bm25_only"]["avg_score"]
        winner = "FAISS" if faiss > bm25 else "BM25"
        print(f"   Retrieval: {winner} wins ({faiss:.3f} vs {bm25:.3f})")
    
    if "full" in results:
        full = results["full"]
        print(f"   Full Pipeline: {full['avg_score']:.3f} [{full['ci_lower']:.3f}, {full['ci_upper']:.3f}]")
        print(f"   Answerable accuracy: {full['answerable_score']:.3f}")
        print(f"   Unanswerable rejection: {full['unanswerable_score']:.3f}")
    
    # Score/Latency efficiency
    print(f"\n   ⚡ SCORE/LATENCY EFFICIENCY:")
    for mode_key in ABLATION_MODES:
        if mode_key in results:
            r = results[mode_key]
            efficiency = r["avg_score"] / max(r["avg_latency"], 0.1)
            print(f"      {r['label']:<16} {efficiency:.4f} score/s")
    
    n = list(results.values())[0]["n_queries"]
    n_ans = list(results.values())[0]["n_answerable"]
    n_unans = list(results.values())[0]["n_unanswerable"]
    print(f"\n   📏 N = {n} queries ({n_ans} answerable + {n_unans} unanswerable)")
    print(f"   📊 {N_BOOTSTRAP} bootstrap iterations, {CONFIDENCE_LEVEL*100:.0f}% CI")
    print(f"   📚 Dataset: SQuAD v1.1 + v2 (English, human ground truth)")


def plot_results(results: Dict):
    """Gráfico completo com CI e Pareto."""
    try:
        import matplotlib.pyplot as plt
        import matplotlib
        matplotlib.rcParams['figure.facecolor'] = '#0d1117'
        matplotlib.rcParams['axes.facecolor'] = '#161b22'
        matplotlib.rcParams['text.color'] = '#c9d1d9'
        matplotlib.rcParams['axes.labelcolor'] = '#c9d1d9'
        matplotlib.rcParams['xtick.color'] = '#8b949e'
        matplotlib.rcParams['ytick.color'] = '#8b949e'
    except ImportError:
        print("⚠️ matplotlib não disponível.")
        return

    labels = [results[k]["label"] for k in ABLATION_MODES if k in results]
    scores = [results[k]["avg_score"] for k in ABLATION_MODES if k in results]
    ci_lower = [results[k]["ci_lower"] for k in ABLATION_MODES if k in results]
    ci_upper = [results[k]["ci_upper"] for k in ABLATION_MODES if k in results]
    latencies = [results[k]["avg_latency"] for k in ABLATION_MODES if k in results]
    ans_scores = [results[k]["answerable_score"] for k in ABLATION_MODES if k in results]
    unans_scores = [results[k]["unanswerable_score"] for k in ABLATION_MODES if k in results]

    colors = ['#6366f1', '#818cf8', '#34d399', '#fbbf24', '#f472b6'][:len(labels)]

    fig, axes = plt.subplots(2, 2, figsize=(16, 12))

    # 1. Scores com CI
    ax1 = axes[0, 0]
    yerr_lower = [s - l for s, l in zip(scores, ci_lower)]
    yerr_upper = [u - s for s, u in zip(scores, ci_upper)]
    bars1 = ax1.bar(labels, scores, color=colors, edgecolor='#30363d', linewidth=1.5,
                    yerr=[yerr_lower, yerr_upper], capsize=5, error_kw={'color': '#f0f6fc', 'linewidth': 2})
    ax1.set_title('Score com 95% CI', fontsize=13, fontweight='bold', pad=10)
    ax1.set_ylabel('Score Médio')
    ax1.set_ylim(0, 1.1)
    ax1.grid(axis='y', alpha=0.15)
    for bar, score in zip(bars1, scores):
        ax1.text(bar.get_x() + bar.get_width()/2., bar.get_height() + 0.04,
                f'{score:.3f}', ha='center', va='bottom', fontweight='bold', fontsize=10)

    # 2. Answerable vs Unanswerable
    ax2 = axes[0, 1]
    x = np.arange(len(labels))
    w = 0.35
    bars_ans = ax2.bar(x - w/2, ans_scores, w, color='#34d399', edgecolor='#30363d', label='Answerable')
    bars_unans = ax2.bar(x + w/2, unans_scores, w, color='#f472b6', edgecolor='#30363d', label='Unanswerable')
    ax2.set_title('Answerable vs Unanswerable', fontsize=13, fontweight='bold', pad=10)
    ax2.set_ylabel('Score')
    ax2.set_ylim(0, 1.15)
    ax2.set_xticks(x)
    ax2.set_xticklabels(labels, rotation=20, ha='right')
    ax2.legend(loc='upper left', framealpha=0.8)
    ax2.grid(axis='y', alpha=0.15)

    # 3. Latência
    ax3 = axes[1, 0]
    lat_stds = [results[k]["std_latency"] for k in ABLATION_MODES if k in results]
    bars3 = ax3.bar(labels, latencies, color=colors, edgecolor='#30363d', linewidth=1.5,
                    yerr=lat_stds, capsize=5, error_kw={'color': '#f0f6fc', 'linewidth': 1.5})
    ax3.set_title('Latência Média (±1σ)', fontsize=13, fontweight='bold', pad=10)
    ax3.set_ylabel('Latência (s)')
    ax3.grid(axis='y', alpha=0.15)
    for bar, lat in zip(bars3, latencies):
        ax3.text(bar.get_x() + bar.get_width()/2., bar.get_height() + 1,
                f'{lat:.1f}s', ha='center', va='bottom', fontweight='bold', fontsize=10)

    # 4. Pareto: Score vs Latência
    ax4 = axes[1, 1]
    for i, (label, score, lat, color) in enumerate(zip(labels, scores, latencies, colors)):
        ax4.scatter(lat, score, color=color, s=200, zorder=5, edgecolors='white', linewidth=2)
        ax4.annotate(label, (lat, score), textcoords="offset points",
                    xytext=(10, 5), fontsize=9, color='#c9d1d9')
    ax4.set_title('Pareto: Score vs Latência', fontsize=13, fontweight='bold', pad=10)
    ax4.set_xlabel('Latência Média (s)')
    ax4.set_ylabel('Score')
    ax4.set_ylim(min(scores) - 0.05, 1.05)
    ax4.grid(alpha=0.15)

    plt.suptitle('AetherMap RAG — Ablation Study v3 (SQuAD EN, N=50)', fontsize=15, fontweight='bold', y=1.01, color='#f0f6fc')
    plt.tight_layout()
    plt.savefig('ablation_v3_results.png', dpi=150, bbox_inches='tight',
                facecolor='#0d1117', edgecolor='none')
    plt.show()
    print("📈 Gráfico salvo em ablation_v3_results.png")


# %%
# ==============================================================================
# EXECUTAR
# ==============================================================================

resultados = run_ablation_study()