VENDOR_NEUTRAL_PROMPT = """You are a clinical natural language processing expert.
Extract 4 fields from the given clinical note as structured JSON only.
Output only a single JSON, no explanation or comments.

Schema:
{
  "symptoms": ["list of strings"],
  "medications": [{"name": "drug name", "dose": "dose", "frequency": "frequency"}],
  "labs": [{"test": "test name", "value": "numeric", "unit": "unit"}],
  "diagnoses": ["list of diagnosis strings"]
}

Principles:
1. Do not infer or invent information not in the note (no hallucination).
2. If a field is absent, empty array [].
3. If JSON parsing fails, the bench is invalidated. JSON only.

Clinical note:
{note}

JSON output:
"""


# Per-vendor optimized prompts (optional bench)
VENDOR_OPTIMIZED_PROMPTS = {
    "anthropic": VENDOR_NEUTRAL_PROMPT,  # Claude is powerful with natural-language prompts
    "openai": VENDOR_NEUTRAL_PROMPT + "\n\n(response_format=json_object)",  # JSON mode hint
    "google": VENDOR_NEUTRAL_PROMPT,  # Gemini uses response_schema separately
    "opensource": VENDOR_NEUTRAL_PROMPT + "\n\nAnswer:",  # Meditron etc. need continuation cue
}

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        # snapshot id (version pinning)
    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:
        """asyncio support. In production use aiohttp · vendor async SDK."""
        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):
    """Local open model (Meditron 7B/70B). Use vLLM · TGI · HF pipeline."""

    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:
    """Lenient JSON parse. Strip markdown fence · surrounding text."""
    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:
        # One more try: fenced JSON block
        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

# Vendor official pricing at execution time (reference, must re-check at article publication time)
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:
    """Tokens × pricing = USD cost."""
    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]:
    """Run one vendor over the entire dataset."""
    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]:
    """Reuse the F1 computation function from Topic 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:
    """Per-vendor 3-axis summary table."""
    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:
    """3-axis scatter plot of accuracy vs cost vs latency + Pareto frontier."""
    import matplotlib.pyplot as plt
    fig, axes = plt.subplots(1, 2, figsize=(16, 6))
    
    # (a) Accuracy vs cost
    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("Accuracy vs Cost (log scale)")
    plt.colorbar(scatter, ax=ax, label="Latency p50 (ms)")
    
    # (b) Accuracy vs latency
    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("Accuracy vs Latency")
    
    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:
    """Extract only Pareto frontier candidates."""
    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:
    """Full vendor bench run · report."""
    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"Running bench: {key}")
        results = await benchmark_vendor(adapter, dataset, VENDOR_NEUTRAL_PROMPT)
        all_results[key] = results
        # Save individual raw results
        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"))
    
    # Extract Pareto frontier
    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"Pareto frontier vendors: {list(pareto['model'])}")
    return report


# Execution example
# 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")))