# ==============================================================================
#  AetherMap — ABLATION STUDY v1
#  Mede o impacto de cada componente do pipeline RAG
#  Para rodar no Google Colab contra a API do HF Space
# ==============================================================================

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

# %%
import os
import json
import time
import requests
import pandas as pd
import numpy as np
import tempfile
import re
from typing import List, Dict, Any

# ==============================================================================
# CONFIGURAÇÃO
# ==============================================================================
OPENROUTER_API_KEY = "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"  # 👈 SUA KEY OPENROUTER
AETHERMAP_URL = "https://madras1-aethermap.hf.space"

# Modelos OpenRouter (gratuitos)
QUERY_GENERATOR_MODEL = "stepfun/step-3.5-flash:free"
EVALUATOR_MODEL = "arcee-ai/trinity-large-preview:free"

DELAY_BETWEEN_LLM_CALLS = 2.0  # Segundos entre chamadas LLM

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

# Modos de ablation e seus componentes
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},
}

print("✅ Configuração carregada")
print(f"🌐 API: {AETHERMAP_URL}")
print(f"🤖 Gerador: {QUERY_GENERATOR_MODEL}")
print(f"🧑‍⚖️ Avaliador: {EVALUATOR_MODEL}")

# %%
# ==============================================================================
# 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  # Reasoning vai pra campo separado, content fica limpo
            },
        }

        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}")
            print(f"   📝 {resp.text[:300]}")
            return ""

        data = resp.json()

        choices = data.get("choices", [])
        if not choices:
            print(f"   ⚠️ LLM ({model}): Sem choices")
            print(f"   📝 Raw: {json.dumps(data)[:400]}")
            return ""

        message = choices[0].get("message", {})
        content = message.get("content") or ""

        if not content:
            print(f"   ⚠️ LLM ({model}): content vazio")
            print(f"   📝 Message keys: {list(message.keys())}")
            print(f"   📝 Raw: {json.dumps(message)[:400]}")
            return ""

        return content.strip()

    except Exception as e:
        print(f"   ⚠️ Erro LLM ({model}): {type(e).__name__}: {e}")
        time.sleep(5)
        return ""


def gerar_queries(docs: List[str], n: int = 10) -> List[Dict]:
    """
    Gera perguntas que podem ser respondidas com os documentos.
    Usa amostra dos docs para gerar queries diversas.
    Com retry e fallback robusto.
    """
    sample_docs = docs[:min(5, len(docs))]
    docs_text = "\n---\n".join([d[:400] for d in sample_docs])

    prompt = f"""Gere {n} perguntas DIVERSAS que podem ser respondidas com estes documentos.
Inclua perguntas fáceis e difíceis. Varie os tipos: factuais, comparativas, de resumo.
Para cada pergunta, dê também a resposta correta baseada nos documentos.

DOCUMENTOS:
{docs_text}

Retorne APENAS JSON válido:
[{{"pergunta": "...", "resposta": "..."}}]"""

    # Tentar até 3 vezes
    for attempt in range(3):
        response = call_llm(prompt, QUERY_GENERATOR_MODEL, max_tokens=4000)

        if not response:
            print(f"   ⚠️ Tentativa {attempt+1}: LLM retornou vazio")
            continue

        try:
            # Encontrar o JSON no response
            start = response.find("[")
            end = response.rfind("]")

            if start == -1:
                print(f"   ⚠️ Tentativa {attempt+1}: Sem '[' no response")
                print(f"   📝 Response: {response[:200]}...")
                continue

            if end == -1 or end <= start:
                # JSON truncado — tentar consertar
                print(f"   ⚠️ Tentativa {attempt+1}: JSON truncado, tentando reparar...")
                json_str = response[start:]
                # Fechar JSON truncado
                if json_str.rstrip().endswith(","):
                    json_str = json_str.rstrip().rstrip(",") + "]"
                elif not json_str.rstrip().endswith("]"):
                    json_str = json_str.rstrip() + '"}]'
            else:
                json_str = response[start:end + 1]

            # Limpar possíveis problemas
            json_str = json_str.replace("'", '"')  # Aspas simples → duplas

            queries = json.loads(json_str)

            if isinstance(queries, list) and len(queries) > 0:
                # Validar que cada item tem as chaves certas
                valid = []
                for q in queries:
                    if isinstance(q, dict):
                        pergunta = q.get("pergunta", q.get("question", ""))
                        resposta = q.get("resposta", q.get("answer", q.get("resposta_esperada", "")))
                        if pergunta and resposta:
                            valid.append({"pergunta": pergunta, "resposta": resposta})

                if valid:
                    print(f"   ✅ {len(valid)} queries extraídas na tentativa {attempt+1}")
                    return valid[:n]

            print(f"   ⚠️ Tentativa {attempt+1}: JSON parseado mas sem queries válidas")

        except json.JSONDecodeError as e:
            print(f"   ⚠️ Tentativa {attempt+1}: JSON inválido — {e}")
            print(f"   📝 Response: {response[:300]}...")

    # Fallback: queries genéricas baseadas nos docs
    print("   🔄 Usando queries fallback genéricas...")
    return [
        {"pergunta": "Qual é o tema principal dos documentos?", "resposta": "Os documentos abordam diversos temas conforme o dataset."},
        {"pergunta": "Quais são as entidades mais mencionadas?", "resposta": "As entidades mais frequentes estão nos documentos."},
        {"pergunta": "Existe alguma tendência nos dados?", "resposta": "Os dados mostram padrões relevantes ao tema."},
        {"pergunta": "Qual a conclusão principal?", "resposta": "A conclusão depende do contexto dos documentos."},
        {"pergunta": "Quais termos técnicos são mais usados?", "resposta": "Os termos técnicos variam conforme o domínio."},
    ][:n]


def avaliar_resposta(pergunta: str, resposta_rag: str, resposta_esperada: str) -> Dict[str, float]:
    """
    Avalia qualidade da resposta do RAG.
    Retorna scores de correção e completude (0-1).
    """
    prompt = f"""Avalie a resposta do sistema comparando com a esperada.

PERGUNTA: {pergunta}
RESPOSTA ESPERADA: {resposta_esperada[:300]}
RESPOSTA DO SISTEMA: {resposta_rag[:400]}

Dê notas de 0.0 a 1.0 para:
- correcao: A informação está factualmente correta?
- completude: A resposta cobre todos os pontos esperados?
- relevancia: A resposta é diretamente relevante à pergunta?

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": float(result.get("correcao", 0.5)),
                    "completude": float(result.get("completude", 0.5)),
                    "relevancia": float(result.get("relevancia", 0.5)),
                }
    except Exception as e:
        print(f"   ⚠️ Erro parsing avaliação: {e}")
        print(f"   📝 Response: {response[:200]}")
    return {"correcao": 0.5, "completude": 0.5, "relevancia": 0.5}


# %%
# ==============================================================================
# UPLOAD DO DATASET
# ==============================================================================

def upload_dataset(file_path: str, text_column: str, n_samples: int = 200) -> str:
    """Faz upload do dataset para o AetherMap e retorna o job_id."""

    # Detectar separador
    with open(file_path, 'r', encoding='utf-8') as f:
        first_line = f.readline()
    sep = ';' if first_line.count(';') > first_line.count(',') else ','

    # Carregar e amostrar
    df = pd.read_csv(file_path, sep=sep)
    sample = df.sample(n=min(n_samples, len(df)), random_state=42)

    # Salvar amostra temporária
    with tempfile.NamedTemporaryFile(mode='w', suffix='.csv', delete=False, encoding='utf-8') as f:
        sample.to_csv(f, index=False)
        temp_path = f.name

    print(f"📤 Upload de {len(sample)} documentos...")

    with open(temp_path, "rb") as f:
        response = requests.post(
            f"{AETHERMAP_URL}/process/",
            files={"file": ("dataset.csv", f)},
            data={
                "n_samples": n_samples,
                "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")

    # Retornar textos para geração de queries
    docs = [p["full_text"] for p in result["plot_data"][:20]]

    return job_id, docs


# %%
# ==============================================================================
# BUSCA COM MODO DE ABLATION
# ==============================================================================

def search_with_mode(job_id: str, query: str, mode: str) -> Dict:
    """Executa busca com um modo de ablation específico."""
    start = time.time()

    response = requests.post(
        f"{AETHERMAP_URL}/search/",
        data={
            "query": query,
            "job_id": job_id,
            "ablation_mode": mode
        },
        timeout=120
    )

    latency = time.time() - start

    if response.status_code != 200:
        return {"error": response.text[:100], "latency": latency}

    result = response.json()
    result["latency"] = latency
    return result


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

def run_ablation_study(
    file_path: str,
    text_column: str = "texto",
    n_samples: int = 200,
    n_queries: int = 10
):
    """
    Executa ablation study completo.

    Args:
        file_path: Caminho do CSV
        text_column: Nome da coluna de texto
        n_samples: Número de documentos para indexar
        n_queries: Número de queries para testar
    """
    print("=" * 60)
    print("🧪 AETHERMAP ABLATION STUDY v1")
    print("=" * 60)

    # 1. Upload
    print("\n📦 FASE 1: Upload do dataset")
    job_id, docs = upload_dataset(file_path, text_column, n_samples)

    # 2. Gerar queries
    print(f"\n🔍 FASE 2: Gerando {n_queries} queries de teste...")
    queries = gerar_queries(docs, n_queries)

    if not queries:
        print("❌ Falha na geração de queries. Abortando.")
        return None

    print(f"   ✅ {len(queries)} queries geradas:")
    for i, q in enumerate(queries):
        print(f"   [{i+1}] {q['pergunta'][:60]}...")

    # 3. Rodar cada modo
    print(f"\n🧪 FASE 3: Testando {len(ABLATION_MODES)} modos de ablation...")

    all_results = {}

    for mode_key, mode_info in ABLATION_MODES.items():
        mode_label = mode_info["label"]
        print(f"\n{'─' * 40}")
        print(f"🔬 Modo: {mode_label} ({mode_key})")
        print(f"   Componentes: ", end="")
        components = []
        if mode_info["faiss"]: components.append("FAISS")
        if mode_info["bm25"]: components.append("BM25")
        if mode_info["rrf"]: components.append("RRF")
        if mode_info["reranker"]: components.append("Reranker")
        if mode_info["expansion"]: components.append("QueryExpansion")
        print(" + ".join(components))

        scores = []
        latencies = []

        for i, q in enumerate(queries):
            pergunta = q.get("pergunta", q.get("question", ""))
            esperada = q.get("resposta", q.get("answer", ""))

            # Buscar
            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))
                continue

            resposta = result.get("summary", "")
            latency = result.get("latency", 0)
            latencies.append(latency)

            # Avaliar
            aval = avaliar_resposta(pergunta, resposta, esperada)
            scores.append(aval)

            avg = (aval["correcao"] + aval["completude"] + aval["relevancia"]) / 3
            emoji = "✓" if avg >= 0.6 else "○"
            print(f"   [{i+1}] {emoji} C:{aval['correcao']:.2f} Com:{aval['completude']:.2f} R:{aval['relevancia']:.2f} | {latency:.1f}s")
            print(f"         Q: {pergunta[:80]}")
            print(f"         RAG: {resposta[:150]}...")
            if avg < 0.3:
                print(f"         ⚠️ Esperado: {esperada[:100]}")

            time.sleep(0.5)  # Rate limiting suave

        # Calcular médias
        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)

        all_results[mode_key] = {
            "label": mode_label,
            "avg_score": avg_score,
            "correcao": avg_correcao,
            "completude": avg_completude,
            "relevancia": avg_relevancia,
            "avg_latency": avg_latency,
            "n_queries": len(queries),
            "components": components,
        }

        print(f"   📊 Score médio: {avg_score:.3f} | Latência: {avg_latency:.1f}s")

    # 4. Resultado Final
    print_results(all_results)
    plot_results(all_results)

    return all_results


# %%
# ==============================================================================
# VISUALIZAÇÃO DOS RESULTADOS
# ==============================================================================

def print_results(results: Dict):
    """Imprime tabela comparativa dos resultados."""
    print("\n" + "=" * 70)
    print("📊 RESULTADO DO ABLATION STUDY")
    print("=" * 70)

    print(f"\n{'Modo':<20} {'Score':>7} {'Correção':>9} {'Complet.':>9} {'Relev.':>7} {'Latência':>9} {'Δ Score':>8}")
    print("─" * 70)

    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]
        delta = r["avg_score"] - baseline_score
        delta_str = f"+{delta:.3f}" if delta > 0 else f"{delta:.3f}"
        trophy = " 🏆" if r["avg_score"] == best_score else ""

        print(f"{r['label']:<20} {r['avg_score']:>7.3f} {r['correcao']:>9.3f} {r['completude']:>9.3f} {r['relevancia']:>7.3f} {r['avg_latency']:>8.1f}s {delta_str:>8}{trophy}")

    # Insights
    print(f"\n{'─' * 70}")
    print("💡 INSIGHTS:")

    if "faiss_only" in results and "hybrid" in results:
        delta = results["hybrid"]["avg_score"] - results["faiss_only"]["avg_score"]
        print(f"   Hybrid Search (RRF): {'+' if delta >= 0 else ''}{delta:.3f} vs FAISS only ({delta/max(results['faiss_only']['avg_score'],0.001)*100:+.1f}%)")

    if "hybrid" in results and "hybrid_rerank" in results:
        delta = results["hybrid_rerank"]["avg_score"] - results["hybrid"]["avg_score"]
        print(f"   + Reranker:          {'+' if delta >= 0 else ''}{delta:.3f} vs Hybrid ({delta/max(results['hybrid']['avg_score'],0.001)*100:+.1f}%)")

    if "hybrid_rerank" in results and "full" in results:
        delta = results["full"]["avg_score"] - results["hybrid_rerank"]["avg_score"]
        print(f"   + Query Expansion:   {'+' if delta >= 0 else ''}{delta:.3f} vs Hybrid+Rerank ({delta/max(results['hybrid_rerank']['avg_score'],0.001)*100:+.1f}%)")


def plot_results(results: Dict):
    """Gera gráfico de barras com os resultados."""
    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. Pulando gráfico.")
        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]
    latencies = [results[k]["avg_latency"] for k in ABLATION_MODES if k in results]

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 6))

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

    # Gráfico 1: Scores
    bars1 = ax1.bar(labels, scores, color=colors, edgecolor='#30363d', linewidth=1.5)
    ax1.set_title('Quality Score por Modo', fontsize=14, fontweight='bold', pad=15)
    ax1.set_ylabel('Score Médio (0-1)')
    ax1.set_ylim(0, 1.05)
    ax1.grid(axis='y', alpha=0.15, color='#8b949e')
    for bar, score in zip(bars1, scores):
        ax1.text(bar.get_x() + bar.get_width()/2., bar.get_height() + 0.02,
                f'{score:.3f}', ha='center', va='bottom', fontweight='bold', fontsize=11)

    # Gráfico 2: Latência
    bars2 = ax2.bar(labels, latencies, color=colors, edgecolor='#30363d', linewidth=1.5)
    ax2.set_title('Latência por Modo', fontsize=14, fontweight='bold', pad=15)
    ax2.set_ylabel('Latência Média (s)')
    ax2.grid(axis='y', alpha=0.15, color='#8b949e')
    for bar, lat in zip(bars2, latencies):
        ax2.text(bar.get_x() + bar.get_width()/2., bar.get_height() + 0.1,
                f'{lat:.1f}s', ha='center', va='bottom', fontweight='bold', fontsize=11)

    plt.xticks(rotation=25, ha='right')
    plt.suptitle('AetherMap RAG — Ablation Study', fontsize=16, fontweight='bold', y=1.02, color='#f0f6fc')
    plt.tight_layout()
    plt.savefig('ablation_results.png', dpi=150, bbox_inches='tight',
                facecolor='#0d1117', edgecolor='none')
    plt.show()
    print("📈 Gráfico salvo em ablation_results.png")


# %%
# ==============================================================================
# EXECUTAR
# ==============================================================================
# Ajuste o caminho do dataset e a coluna de texto:

resultados = run_ablation_study(
    file_path="dataset_unificado_72k.csv",
    text_column="texto",
    n_samples=200,
    n_queries=20
)