VENDOR_NEUTRAL_PROMPT = """Eres un experto en procesamiento del lenguaje natural clínico.
Extrae únicamente los 4 campos en JSON estructurado a partir de la nota clínica proporcionada.
Imprime solo un JSON sin explicaciones ni comentarios.

Esquema:
{
  "symptoms": ["lista de cadenas"],
  "medications": [{"name": "nombre del medicamento", "dose": "dosis", "frequency": "frecuencia"}],
  "labs": [{"test": "nombre de la prueba", "value": "valor", "unit": "unidad"}],
  "diagnoses": ["lista de cadenas con nombres de diagnóstico"]
}

Principios:
1. No infieras ni crees información que no esté en la nota (prohibido alucinar).
2. Si falta un campo, usa un array vacío [].
3. Si falla el análisis JSON, el benchmark se invalida. Solo JSON.

Nota clínica:
{note}

Salida JSON:
"""


# Prompts especializados por proveedor (benchmark opcional)
VENDOR_OPTIMIZED_PROMPTS = {
    "anthropic": VENDOR_NEUTRAL_PROMPT,  # Claude es potente por sí mismo en prompts de lenguaje natural
    "openai": VENDOR_NEUTRAL_PROMPT + "\n\n(response_format=json_object)",  # Pista del modo JSON
    "google": VENDOR_NEUTRAL_PROMPT,  # Gemini requiere especificar response_schema por separado
    "opensource": VENDOR_NEUTRAL_PROMPT + "\n\nRespuesta:",  # Meditron y otros requieren continuación explícita
}

import asyncio
import time
import json
import re
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Any

import anthropic
from openai import OpenAI
import google.generativeai as genai


@dataclass
class VendorResponse:
    vendor: str
    model: str
    model_snapshot: str        # id de la instantánea (fijación de versión)
    parsed_output: dict | None
    raw_text: str
    input_tokens: int
    output_tokens: int
    latency_ms: float
    error: str | None
    retry_count: int = 0


class VendorAdapter(ABC):
    def __init__(self, model: str, max_retries: int = 3):
        self.model = model
        self.max_retries = max_retries

    @abstractmethod
    def query(self, prompt: str) -> VendorResponse:
        ...

    async def query_async(self, prompt: str) -> VendorResponse:
        """Compatible con asyncio. En producción usar aiohttp · SDK asíncrono del proveedor."""
        return await asyncio.get_event_loop().run_in_executor(None, self.query, prompt)


class ClaudeAdapter(VendorAdapter):
    def __init__(self, model: str = "claude-sonnet-4-5-20241022"):
        super().__init__(model)
        self.client = anthropic.Anthropic()
        self.snapshot = model

    def query(self, prompt: str) -> VendorResponse:
        start = time.perf_counter()
        retry = 0
        for attempt in range(self.max_retries):
            try:
                resp = self.client.messages.create(
                    model=self.model,
                    max_tokens=2048,
                    messages=[{"role": "user", "content": prompt}],
                )
                elapsed = (time.perf_counter() - start) * 1000
                raw = resp.content[0].text
                parsed = try_parse_json(raw)
                return VendorResponse(
                    vendor="anthropic", model=self.model, model_snapshot=self.snapshot,
                    parsed_output=parsed, raw_text=raw,
                    input_tokens=resp.usage.input_tokens,
                    output_tokens=resp.usage.output_tokens,
                    latency_ms=elapsed, error=None, retry_count=retry,
                )
            except anthropic.RateLimitError:
                retry += 1
                time.sleep(2 ** attempt)
            except Exception as e:
                return VendorResponse(
                    vendor="anthropic", model=self.model, model_snapshot=self.snapshot,
                    parsed_output=None, raw_text="",
                    input_tokens=0, output_tokens=0,
                    latency_ms=(time.perf_counter() - start) * 1000,
                    error=str(e), retry_count=retry,
                )
        return VendorResponse(
            vendor="anthropic", model=self.model, model_snapshot=self.snapshot,
            parsed_output=None, raw_text="",
            input_tokens=0, output_tokens=0,
            latency_ms=(time.perf_counter() - start) * 1000,
            error="max_retries_exceeded", retry_count=retry,
        )


class OpenAIAdapter(VendorAdapter):
    def __init__(self, model: str = "gpt-4o-2024-08-06", use_json_mode: bool = True):
        super().__init__(model)
        self.client = OpenAI()
        self.snapshot = model
        self.use_json_mode = use_json_mode

    def query(self, prompt: str) -> VendorResponse:
        start = time.perf_counter()
        retry = 0
        for attempt in range(self.max_retries):
            try:
                kwargs = {
                    "model": self.model,
                    "messages": [{"role": "user", "content": prompt}],
                    "max_tokens": 2048,
                }
                if self.use_json_mode and not self.model.startswith("o"):
                    kwargs["response_format"] = {"type": "json_object"}
                resp = self.client.chat.completions.create(**kwargs)
                elapsed = (time.perf_counter() - start) * 1000
                raw = resp.choices[0].message.content or ""
                parsed = try_parse_json(raw)
                return VendorResponse(
                    vendor="openai", model=self.model, model_snapshot=self.snapshot,
                    parsed_output=parsed, raw_text=raw,
                    input_tokens=resp.usage.prompt_tokens,
                    output_tokens=resp.usage.completion_tokens,
                    latency_ms=elapsed, error=None, retry_count=retry,
                )
            except Exception as e:
                if "rate_limit" in str(e).lower():
                    retry += 1
                    time.sleep(2 ** attempt)
                    continue
                return VendorResponse(
                    vendor="openai", model=self.model, model_snapshot=self.snapshot,
                    parsed_output=None, raw_text="",
                    input_tokens=0, output_tokens=0,
                    latency_ms=(time.perf_counter() - start) * 1000,
                    error=str(e), retry_count=retry,
                )
        return VendorResponse(
            vendor="openai", model=self.model, model_snapshot=self.snapshot,
            parsed_output=None, raw_text="",
            input_tokens=0, output_tokens=0,
            latency_ms=(time.perf_counter() - start) * 1000,
            error="max_retries_exceeded", retry_count=retry,
        )


class GeminiAdapter(VendorAdapter):
    def __init__(self, model: str = "gemini-2.5-pro"):
        super().__init__(model)
        genai.configure()
        self.snapshot = model
        self.model_obj = genai.GenerativeModel(model)

    def query(self, prompt: str) -> VendorResponse:
        start = time.perf_counter()
        try:
            resp = self.model_obj.generate_content(prompt)
            elapsed = (time.perf_counter() - start) * 1000
            raw = resp.text if hasattr(resp, "text") else ""
            parsed = try_parse_json(raw)
            usage = getattr(resp, "usage_metadata", None)
            input_toks = getattr(usage, "prompt_token_count", 0) if usage else 0
            output_toks = getattr(usage, "candidates_token_count", 0) if usage else 0
            return VendorResponse(
                vendor="google", model=self.model, model_snapshot=self.snapshot,
                parsed_output=parsed, raw_text=raw,
                input_tokens=input_toks, output_tokens=output_toks,
                latency_ms=elapsed, error=None,
            )
        except Exception as e:
            return VendorResponse(
                vendor="google", model=self.model, model_snapshot=self.snapshot,
                parsed_output=None, raw_text="",
                input_tokens=0, output_tokens=0,
                latency_ms=(time.perf_counter() - start) * 1000,
                error=str(e),
            )


class MeditronAdapter(VendorAdapter):
    """Modelos locales de código abierto (Meditron 7B/70B). Utiliza vLLM · TGI · pipeline de HF."""

    def __init__(self, model_id: str = "epfl-llm/meditron-7b", device: str = "cuda"):
        super().__init__(model_id)
        from transformers import AutoTokenizer, AutoModelForCausalLM
        import torch
        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
        self.model_obj = AutoModelForCausalLM.from_pretrained(
            model_id, torch_dtype=torch.float16, device_map=device,
        )
        self.device = device
        self.snapshot = model_id
        self.torch = torch

    def query(self, prompt: str) -> VendorResponse:
        start = time.perf_counter()
        try:
            inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
            with self.torch.no_grad():
                outputs = self.model_obj.generate(
                    **inputs, max_new_tokens=2048, do_sample=False, temperature=0.1,
                )
            elapsed = (time.perf_counter() - start) * 1000
            raw = self.tokenizer.decode(
                outputs[0][inputs.input_ids.size(1):], skip_special_tokens=True,
            )
            parsed = try_parse_json(raw)
            return VendorResponse(
                vendor="opensource", model=self.model, model_snapshot=self.snapshot,
                parsed_output=parsed, raw_text=raw,
                input_tokens=int(inputs.input_ids.size(1)),
                output_tokens=int(outputs.size(1) - inputs.input_ids.size(1)),
                latency_ms=elapsed, error=None,
            )
        except Exception as e:
            return VendorResponse(
                vendor="opensource", model=self.model, model_snapshot=self.snapshot,
                parsed_output=None, raw_text="",
                input_tokens=0, output_tokens=0,
                latency_ms=(time.perf_counter() - start) * 1000,
                error=str(e),
            )


def try_parse_json(text: str) -> dict | None:
    """Análisis JSON indulgente. Elimina la cerca de markdown y el texto circundante."""
    if not text:
        return None
    match = re.search(r"\{.*\}", text, re.DOTALL)
    if not match:
        return None
    try:
        return json.loads(match.group(0))
    except json.JSONDecodeError:
        # Intento adicional: bloque JSON con cercas
        fence = re.search(r"```(?:json)?\s*(\{.*?\})\s*```", text, re.DOTALL)
        if fence:
            try:
                return json.loads(fence.group(1))
            except json.JSONDecodeError:
                pass
    return None

# Tarifas oficiales del proveedor en tiempo de ejecución (referencia, verificar al momento de la publicación)
PRICING_TABLE_USD_PER_M_TOKENS = {
    "claude-opus-4-5-20250219": {"input": 15.0, "output": 75.0},
    "claude-sonnet-4-5-20241022": {"input": 3.0, "output": 15.0},
    "claude-haiku-4-5-20251001": {"input": 1.0, "output": 5.0},
    "gpt-4o-2024-08-06": {"input": 2.5, "output": 10.0},
    "gpt-4o-mini-2024-07-18": {"input": 0.15, "output": 0.60},
    "o1-2024-12-17": {"input": 15.0, "output": 60.0},
    "o3-2025-04-16": {"input": 10.0, "output": 40.0},
    "gemini-2.5-pro": {"input": 1.25, "output": 5.0},
    "gemini-2.5-flash": {"input": 0.15, "output": 0.6},
    "epfl-llm/meditron-7b": {"input": 0.0, "output": 0.0},  # Local
    "epfl-llm/meditron-70b": {"input": 0.0, "output": 0.0},
}


def calculate_cost_usd(model_snapshot: str, input_tokens: int, output_tokens: int) -> float:
    """Costo en USD = número de tokens × tarifa."""
    pricing = PRICING_TABLE_USD_PER_M_TOKENS.get(model_snapshot, {"input": 0.0, "output": 0.0})
    return (input_tokens * pricing["input"] + output_tokens * pricing["output"]) / 1_000_000

from dataclasses import dataclass


@dataclass
class BenchResult:
    vendor_response: VendorResponse
    cost_usd: float
    f1_scores: dict[str, float]
    gold_label: dict


async def benchmark_vendor(
    adapter: VendorAdapter,
    dataset: list[dict],  # [{note, gold_label}]
    prompt_template: str = VENDOR_NEUTRAL_PROMPT,
    max_concurrent: int = 5,
) -> list[BenchResult]:
    """Ejecutar el conjunto de datos completo para un proveedor."""
    semaphore = asyncio.Semaphore(max_concurrent)

    async def process_one(sample: dict) -> BenchResult:
        async with semaphore:
            prompt = prompt_template.replace("{note}", sample["note"])
            resp = await adapter.query_async(prompt)
            cost = calculate_cost_usd(resp.model_snapshot, resp.input_tokens, resp.output_tokens)
            if resp.parsed_output:
                f1 = evaluate_extraction(resp.parsed_output, sample["gold_label"])
            else:
                f1 = {"symptoms": 0.0, "medications": 0.0, "labs": 0.0, "diagnoses": 0.0}
            return BenchResult(
                vendor_response=resp, cost_usd=cost,
                f1_scores=f1, gold_label=sample["gold_label"],
            )

    return await asyncio.gather(*[process_one(s) for s in dataset])


def evaluate_extraction(pred: dict, gold: dict) -> dict[str, float]:
    """Reutilizar la función de cálculo F1 del lote 01."""
    scores = {}
    for field in ["symptoms", "diagnoses"]:
        p = {s.lower().strip() for s in pred.get(field, []) if isinstance(s, str)}
        g = {s.lower().strip() for s in gold.get(field, []) if isinstance(s, str)}
        scores[field] = f1_score(p, g)
    for field in ["medications", "labs"]:
        p = {json.dumps(x, sort_keys=True) for x in pred.get(field, []) if isinstance(x, dict)}
        g = {json.dumps(x, sort_keys=True) for x in gold.get(field, []) if isinstance(x, dict)}
        scores[field] = f1_score(p, g)
    return scores


def f1_score(pred: set, gold: set) -> float:
    if not pred and not gold:
        return 1.0
    tp = len(pred & gold)
    if tp == 0:
        return 0.0
    precision = tp / len(pred)
    recall = tp / len(gold)
    return 2 * precision * recall / (precision + recall)

import numpy as np
import pandas as pd


def aggregate_report(all_results: dict[str, list[BenchResult]]) -> pd.DataFrame:
    """Tabla resumen de tres ejes por proveedor."""
    rows = []
    for vendor_key, results in all_results.items():
        valid = [r for r in results if r.vendor_response.error is None and r.vendor_response.parsed_output]
        n_valid = len(valid)
        n_total = len(results)
        if n_valid == 0:
            rows.append({
                "vendor": vendor_key, "n_valid": 0, "n_total": n_total,
                "macro_f1": 0.0, "cost_per_1k_usd": 0.0,
                "latency_p50_ms": 0.0, "latency_p95_ms": 0.0, "latency_p99_ms": 0.0,
                "failure_rate": 1.0,
            })
            continue

        avg_f1 = {
            field: float(np.mean([r.f1_scores[field] for r in valid]))
            for field in ["symptoms", "medications", "labs", "diagnoses"]
        }
        macro_f1 = float(np.mean(list(avg_f1.values())))

        total_cost = sum(r.cost_usd for r in valid)
        cost_per_1k = (total_cost / n_valid) * 1000

        latencies = [r.vendor_response.latency_ms for r in valid]
        p50 = float(np.percentile(latencies, 50))
        p95 = float(np.percentile(latencies, 95))
        p99 = float(np.percentile(latencies, 99))

        failure_rate = 1 - (n_valid / n_total)

        rows.append({
            "vendor": vendor_key,
            "model": valid[0].vendor_response.model,
            "snapshot": valid[0].vendor_response.model_snapshot,
            "n_valid": n_valid, "n_total": n_total,
            "macro_f1": round(macro_f1, 4),
            **{f"f1_{k}": round(v, 4) for k, v in avg_f1.items()},
            "cost_per_1k_usd": round(cost_per_1k, 3),
            "latency_p50_ms": round(p50, 1),
            "latency_p95_ms": round(p95, 1),
            "latency_p99_ms": round(p99, 1),
            "failure_rate": round(failure_rate, 4),
        })
    return pd.DataFrame(rows).sort_values("macro_f1", ascending=False)


def plot_3axis_pareto(report_df: pd.DataFrame, output_path: str = "bench_pareto.png") -> None:
    """Gráfico de dispersión de tres ejes: precisión vs costo vs latencia + frontera de Pareto."""
    import matplotlib.pyplot as plt
    fig, axes = plt.subplots(1, 2, figsize=(16, 6))

    # (a) Precisión vs costo
    ax = axes[0]
    scatter = ax.scatter(
        report_df["cost_per_1k_usd"], report_df["macro_f1"],
        c=report_df["latency_p50_ms"], s=200, cmap="viridis",
    )
    for _, row in report_df.iterrows():
        ax.annotate(row["model"], (row["cost_per_1k_usd"], row["macro_f1"]),
                    xytext=(5, 5), textcoords="offset points", fontsize=8)
    ax.set_xlabel("Cost per 1000 requests (USD)")
    ax.set_ylabel("Macro F1")
    ax.set_xscale("log")
    ax.set_title("Precisión vs costo (escala log)")
    plt.colorbar(scatter, ax=ax, label="Latency p50 (ms)")

    # (b) Precisión vs latencia
    ax2 = axes[1]
    ax2.scatter(
        report_df["latency_p50_ms"], report_df["macro_f1"],
        s=200, c="steelblue",
    )
    for _, row in report_df.iterrows():
        ax2.annotate(row["model"], (row["latency_p50_ms"], row["macro_f1"]),
                     xytext=(5, 5), textcoords="offset points", fontsize=8)
    ax2.set_xlabel("Latency p50 (ms)")
    ax2.set_ylabel("Macro F1")
    ax2.set_title("Precisión vs latencia")

    plt.tight_layout()
    plt.savefig(output_path, dpi=150, bbox_inches="tight")


def identify_pareto_frontier(df: pd.DataFrame, higher_better: list[str], lower_better: list[str]) -> pd.DataFrame:
    """Extraer solo los candidatos de la frontera de Pareto."""
    pareto = df.copy()
    is_pareto = np.ones(len(df), dtype=bool)
    for i in range(len(df)):
        for j in range(len(df)):
            if i == j:
                continue
            better_or_equal_all = True
            strictly_better_any = False
            for col in higher_better:
                if df.iloc[j][col] < df.iloc[i][col]:
                    better_or_equal_all = False
                    break
                if df.iloc[j][col] > df.iloc[i][col]:
                    strictly_better_any = True
            if not better_or_equal_all:
                continue
            for col in lower_better:
                if df.iloc[j][col] > df.iloc[i][col]:
                    better_or_equal_all = False
                    break
                if df.iloc[j][col] < df.iloc[i][col]:
                    strictly_better_any = True
            if better_or_equal_all and strictly_better_any:
                is_pareto[i] = False
                break
    pareto = pareto[is_pareto]
    return pareto

from pathlib import Path


async def full_benchmark(
    dataset: list[dict],
    vendors: list[VendorAdapter],
    output_dir: Path,
) -> pd.DataFrame:
    """Ejecución y reporte completo del benchmark por proveedor."""
    output_dir.mkdir(parents=True, exist_ok=True)
    all_results = {}
    for adapter in vendors:
        key = f"{adapter.__class__.__name__}_{adapter.model.split('/')[-1]}"
        print(f"Benchmark en ejecución: {key}")
        results = await benchmark_vendor(adapter, dataset, VENDOR_NEUTRAL_PROMPT)
        all_results[key] = results
        # Guardar resultados raw individuales
        with open(output_dir / f"raw_{key}.jsonl", "w") as f:
            for r in results:
                f.write(json.dumps({
                    "model": r.vendor_response.model,
                    "parsed": r.vendor_response.parsed_output,
                    "f1": r.f1_scores,
                    "latency_ms": r.vendor_response.latency_ms,
                    "cost_usd": r.cost_usd,
                    "error": r.vendor_response.error,
                }) + "\n")

    report = aggregate_report(all_results)
    report.to_csv(output_dir / "bench_report.csv", index=False)
    plot_3axis_pareto(report, str(output_dir / "bench_pareto.png"))

    # Extraer frontera de Pareto
    pareto = identify_pareto_frontier(
        report, higher_better=["macro_f1"], lower_better=["cost_per_1k_usd", "latency_p50_ms"],
    )
    pareto.to_csv(output_dir / "pareto_frontier.csv", index=False)
    print(f"Vendedores de la frontera de Pareto: {list(pareto['model'])}")
    return report


# Ejemplo de ejecución
# vendors = [
#     ClaudeAdapter("claude-sonnet-4-5-20241022"),
#     ClaudeAdapter("claude-opus-4-5-20250219"),
#     OpenAIAdapter("gpt-4o-2024-08-06"),
#     OpenAIAdapter("gpt-4o-mini-2024-07-18"),
#     GeminiAdapter("gemini-2.5-pro"),
#     GeminiAdapter("gemini-2.5-flash"),
#     MeditronAdapter("epfl-llm/meditron-7b"),
# ]
# report = asyncio.run(full_benchmark(dataset, vendors, Path("./bench_output")))