VENDOR_NEUTRAL_PROMPT = """당신은 임상 자연어 처리 전문가입니다.
주어진 임상 노트에서 4개 필드를 정형 JSON으로만 추출하세요.
설명·주석 없이 JSON 하나만 출력하세요.

스키마:
{
  "symptoms": ["문자열 리스트"],
  "medications": [{"name": "약명", "dose": "용량", "frequency": "빈도"}],
  "labs": [{"test": "검사명", "value": "수치", "unit": "단위"}],
  "diagnoses": ["진단명 문자열 리스트"]
}

원칙:
1. 노트에 없는 정보를 추론·창작하지 않는다 (할루시네이션 금지).
2. 필드가 없으면 빈 배열 [].
3. JSON 파싱 실패 시 벤치가 무효 처리됨. 반드시 JSON만.

임상 노트:
{note}

JSON 출력:
"""


# 벤더별 특화 프롬프트 (선택적 벤치)
VENDOR_OPTIMIZED_PROMPTS = {
    "anthropic": VENDOR_NEUTRAL_PROMPT,  # Claude는 자연어 프롬프트 자체가 강력
    "openai": VENDOR_NEUTRAL_PROMPT + "\n\n(response_format=json_object)",  # JSON mode 힌트
    "google": VENDOR_NEUTRAL_PROMPT,  # Gemini는 response_schema 별도 지정
    "opensource": VENDOR_NEUTRAL_PROMPT + "\n\n답변:",  # Meditron 등은 continuation 명시
}

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 (버전 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 지원. 실전은 aiohttp · 벤더 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):
    """로컬 오픈 모델 (Meditron 7B/70B). 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:
    """관대한 JSON 파싱. markdown fence · 주변 텍스트 제거."""
    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:
        # 하나 더 시도: fenced JSON 블록
        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

# 실행 시점 벤더 공식 요금표 (참고, 편 발행 시점 재확인 필수)
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},  # 로컬
    "epfl-llm/meditron-70b": {"input": 0.0, "output": 0.0},
}


def calculate_cost_usd(model_snapshot: str, input_tokens: int, output_tokens: int) -> float:
    """토큰 수 × 요금표 = USD 비용."""
    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]:
    """한 벤더에 대해 데이터셋 전체 실행."""
    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]:
    """편 01의 F1 계산 함수 재사용."""
    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:
    """벤더별 3축 요약 표."""
    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:
    """정확도 vs 비용 vs 지연 3축 산점도 + Pareto frontier."""
    import matplotlib.pyplot as plt
    fig, axes = plt.subplots(1, 2, figsize=(16, 6))
    
    # (a) 정확도 vs 비용
    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("정확도 vs 비용 (log scale)")
    plt.colorbar(scatter, ax=ax, label="Latency p50 (ms)")
    
    # (b) 정확도 vs 지연
    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("정확도 vs 지연")
    
    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:
    """Pareto frontier 후보만 추출."""
    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:
    """전체 벤더 벤치 실행 · 리포트."""
    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"벤치 실행: {key}")
        results = await benchmark_vendor(adapter, dataset, VENDOR_NEUTRAL_PROMPT)
        all_results[key] = results
        # 개별 raw 결과 저장
        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"))
    
    # 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 벤더: {list(pareto['model'])}")
    return report


# 실행 예시
# 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")))