from dataclasses import dataclass, field
from typing import Literal

from datasets import load_dataset


@dataclass
class BenchQuestion:
    bench_name: str
    qid: str
    question: str
    choices: list[str] | None       # 객관식이면 리스트, 자유서술이면 None
    gold_answer: str                 # 객관식은 "A"/"B"/..., 자유서술은 rubric JSON
    metadata: dict = field(default_factory=dict)


def load_medqa(split: str = "test", max_samples: int | None = None) -> list[BenchQuestion]:
    """MedQA (USMLE style) 로딩. bigbio/med_qa 리포지토리."""
    ds = load_dataset("bigbio/med_qa", "med_qa_en_source", split=split)
    if max_samples:
        ds = ds.select(range(min(max_samples, len(ds))))
    questions = []
    for i, row in enumerate(ds):
        questions.append(BenchQuestion(
            bench_name="medqa",
            qid=f"medqa_{i}",
            question=row["question"],
            choices=[opt["value"] for opt in row["options"]],
            gold_answer=row["answer_idx"],  # "A", "B", ...
            metadata={"meta_info": row.get("meta_info", "")},
        ))
    return questions


def load_pubmedqa(split: str = "train", max_samples: int | None = None) -> list[BenchQuestion]:
    """PubMedQA (yes/no/maybe). qiaojin/PubMedQA labeled subset."""
    ds = load_dataset("qiaojin/PubMedQA", "pqa_labeled", split=split)
    if max_samples:
        ds = ds.select(range(min(max_samples, len(ds))))
    questions = []
    for i, row in enumerate(ds):
        context = " ".join(row["context"]["contexts"])
        questions.append(BenchQuestion(
            bench_name="pubmedqa",
            qid=f"pmqa_{i}",
            question=f"컨텍스트: {context}\n\n질문: {row['question']}\n답변은 yes/no/maybe 중 하나:",
            choices=["yes", "no", "maybe"],
            gold_answer=row["final_decision"],
            metadata={"pmid": row.get("pubid", "")},
        ))
    return questions


def load_mmlu_medical(split: str = "test", max_samples: int | None = None) -> list[BenchQuestion]:
    """MMLU의 의료 subset. cais/mmlu."""
    subsets = ["clinical_knowledge", "college_medicine", "anatomy", "professional_medicine", "medical_genetics"]
    questions = []
    for subset in subsets:
        ds = load_dataset("cais/mmlu", subset, split=split)
        if max_samples:
            ds = ds.select(range(min(max_samples // len(subsets), len(ds))))
        for i, row in enumerate(ds):
            questions.append(BenchQuestion(
                bench_name=f"mmlu_{subset}",
                qid=f"mmlu_{subset}_{i}",
                question=row["question"],
                choices=row["choices"],
                gold_answer=["A", "B", "C", "D"][row["answer"]],
                metadata={"subset": subset},
            ))
    return questions


def load_medmcqa(split: str = "validation", max_samples: int | None = None) -> list[BenchQuestion]:
    """MedMCQA (인도 의사고시). openlifescienceai/medmcqa."""
    try:
        ds = load_dataset("openlifescienceai/medmcqa", split=split)
    except Exception:
        return []
    if max_samples:
        ds = ds.select(range(min(max_samples, len(ds))))
    questions = []
    for i, row in enumerate(ds):
        questions.append(BenchQuestion(
            bench_name="medmcqa",
            qid=f"medmcqa_{i}",
            question=row["question"],
            choices=[row["opa"], row["opb"], row["opc"], row["opd"]],
            gold_answer=["A", "B", "C", "D"][row["cop"]],
            metadata={"subject_name": row.get("subject_name", "")},
        ))
    return questions


def load_healthbench(split: str = "test", max_samples: int | None = None) -> list[BenchQuestion]:
    """HealthBench (OpenAI 2025). 저장소·라이선스 사전 확인 필요.
    
    실제 저장소 이름은 발표 시점 재확인. 여기서는 개념 스텁.
    """
    try:
        ds = load_dataset("openai/healthbench", split=split)
    except Exception:
        # 대체: 논문 부록 sample scenarios 수동 로딩
        return []
    if max_samples:
        ds = ds.select(range(min(max_samples, len(ds))))
    questions = []
    for i, row in enumerate(ds):
        questions.append(BenchQuestion(
            bench_name="healthbench",
            qid=f"hb_{i}",
            question=row["scenario"],
            choices=None,
            gold_answer=row.get("rubric_json", "{}"),
            metadata=row.get("metadata", {}),
        ))
    return questions

ZERO_SHOT_MEDQA = """다음 의사 국가고시 문제에 답하세요.

문제: {question}

선택지:
{choices_formatted}

답 (A/B/C/D/E 한 글자만):
"""


ZERO_SHOT_COT_MEDQA = """다음 의사 국가고시 문제에 답하세요. 먼저 단계별로 추론한 후 최종 답을 제시하세요.

문제: {question}

선택지:
{choices_formatted}

단계별 추론:
"""


FEW_SHOT_MEDQA = """다음은 의사 국가고시 문제와 정답 예시입니다.

예시 1:
문제: 65세 남성 환자, 3시간 전 갑작스러운 흉통 · 발한 · 오심 호소. 심전도상 ST 상승. 가장 가능성 높은 진단은?
선택지: A. 급성 심근경색 B. 불안정 협심증 C. 대동맥 박리 D. 폐색전증 E. 심낭염
답: A

예시 2:
문제: 45세 여성, 야간 발한 · 체중 감소 · 만성 기침 3개월. 흉부 X선상 우상엽 결절. 가장 우선 시행할 검사는?
선택지: A. CT 스캔 B. 객담 항산균 검사 C. 결핵균 배양 D. 기관지 내시경 E. Mantoux 검사
답: B

이제 다음 문제에 답하세요.

문제: {question}

선택지:
{choices_formatted}

답 (한 글자만):
"""


ZERO_SHOT_PUBMEDQA = """{question}
"""


ZERO_SHOT_MMLU = """다음 문제에 답하세요.

문제: {question}

선택지:
A. {a}
B. {b}
C. {c}
D. {d}

답 (A/B/C/D 한 글자만):
"""


def format_question(q: BenchQuestion, prompt_type: str = "zero_shot") -> str:
    """벤치별 프롬프트 렌더링."""
    if q.bench_name == "medqa":
        choices_text = "\n".join(f"{chr(65+i)}. {c}" for i, c in enumerate(q.choices or []))
        template = {
            "zero_shot": ZERO_SHOT_MEDQA,
            "zero_shot_cot": ZERO_SHOT_COT_MEDQA,
            "few_shot": FEW_SHOT_MEDQA,
        }.get(prompt_type, ZERO_SHOT_MEDQA)
        return template.format(question=q.question, choices_formatted=choices_text)
    
    elif q.bench_name == "pubmedqa":
        return ZERO_SHOT_PUBMEDQA.format(question=q.question)
    
    elif q.bench_name.startswith("mmlu_"):
        return ZERO_SHOT_MMLU.format(
            question=q.question,
            a=q.choices[0], b=q.choices[1], c=q.choices[2], d=q.choices[3],
        )
    
    elif q.bench_name == "medmcqa":
        choices_text = "\n".join(f"{chr(65+i)}. {c}" for i, c in enumerate(q.choices or []))
        return ZERO_SHOT_MEDQA.format(question=q.question, choices_formatted=choices_text)
    
    elif q.bench_name == "healthbench":
        return q.question  # 자유서술
    
    else:
        raise ValueError(f"Unknown bench: {q.bench_name}")

import re


def extract_choice(response: str, choices: list[str] | None) -> str | None:
    """LLM 응답에서 A/B/C/D 하나 추출.
    
    CoT 응답 (긴 추론 뒤 답)도 처리.
    """
    if not choices:
        return None
    upper = response.upper()
    # "답: A" · "정답은 A" · "The answer is A" 같은 패턴 우선
    for pattern in [
        r"답[:\s]+([A-E])",
        r"정답[은:\s]+([A-E])",
        r"answer[:\s]+([A-E])",
        r"the answer is\s+([A-E])",
        r"final answer[:\s]+([A-E])",
    ]:
        m = re.search(pattern, upper)
        if m:
            return m.group(1)
    # Fallback: 마지막에 등장하는 단독 A~E
    matches = re.findall(r"\b([A-E])\b", upper)
    if matches:
        return matches[-1]
    # Fallback: 선택지 텍스트 매칭
    for i, choice in enumerate(choices):
        if choice and choice.lower() in response.lower():
            return chr(65 + i)
    return None


def extract_yes_no_maybe(response: str) -> str | None:
    """PubMedQA yes/no/maybe 추출."""
    text = response.lower().strip()
    for keyword in ["yes", "no", "maybe"]:
        if re.search(rf"\b{keyword}\b", text[:100]):
            return keyword
    return None

JUDGE_PROMPT = """당신은 임상 진단 · 치료 평가 전문가입니다.
아래 시나리오와 응답을 rubric에 따라 평가하세요.

시나리오:
{scenario}

응답:
{response}

Rubric (JSON, 각 항목이 0 or 1로 평가):
{rubric}

각 rubric 항목에 대해 0(미충족) · 1(충족) 점수를 JSON으로 출력:
{{
  "criterion_1": 0 or 1,
  "criterion_2": 0 or 1,
  ...
  "total_score": 합계,
  "max_score": rubric 총 항목 수
}}
"""


def judge_healthbench_response(
    scenario: str,
    response: str,
    rubric_json: str,
    judge_client,
    judge_model: str = "claude-opus-4-5-20250219",
) -> dict:
    """LLM-as-judge로 rubric 스코어링."""
    prompt = JUDGE_PROMPT.format(scenario=scenario, response=response, rubric=rubric_json)
    resp = judge_client.messages.create(
        model=judge_model,
        max_tokens=1024,
        messages=[{"role": "user", "content": prompt}],
    )
    raw = resp.content[0].text
    match = re.search(r"\{.*\}", raw, re.DOTALL)
    if not match:
        return {"total_score": 0, "max_score": 1, "error": "판정 파싱 실패"}
    import json
    try:
        return json.loads(match.group(0))
    except json.JSONDecodeError:
        return {"total_score": 0, "max_score": 1, "error": "JSON invalid"}

import asyncio
from dataclasses import dataclass


@dataclass
class BenchAnswerResult:
    question: BenchQuestion
    model: str
    model_snapshot: str
    response: str
    predicted_answer: str | None
    is_correct: bool
    latency_ms: float
    metadata: dict


async def run_bench_on_vendor(
    questions: list[BenchQuestion],
    adapter,  # 편 09의 VendorAdapter
    prompt_type: str = "zero_shot",
    max_concurrent: int = 5,
) -> list[BenchAnswerResult]:
    """벤치 하나 × 모델 하나 실행."""
    semaphore = asyncio.Semaphore(max_concurrent)
    
    async def process(q: BenchQuestion) -> BenchAnswerResult:
        async with semaphore:
            prompt = format_question(q, prompt_type=prompt_type)
            resp = await adapter.query_async(prompt)
            if q.choices:
                if q.bench_name == "pubmedqa":
                    predicted = extract_yes_no_maybe(resp.raw_text)
                else:
                    predicted = extract_choice(resp.raw_text, q.choices)
                is_correct = predicted == q.gold_answer
            else:
                # HealthBench 등 자유서술은 별도 judge
                predicted = None
                is_correct = False
            return BenchAnswerResult(
                question=q, model=adapter.model, model_snapshot=resp.model_snapshot,
                response=resp.raw_text, predicted_answer=predicted,
                is_correct=is_correct, latency_ms=resp.latency_ms,
                metadata={
                    "input_tokens": resp.input_tokens,
                    "output_tokens": resp.output_tokens,
                    "prompt_type": prompt_type,
                },
            )
    
    return await asyncio.gather(*[process(q) for q in questions])

import numpy as np
import pandas as pd


# 벤더 · 논문 발표 숫자 (실행 시점 재확인 필수)
PUBLISHED_NUMBERS = {
    "medqa": {
        "claude-opus-4-5": 0.905,      # 예시, Anthropic 발표
        "gpt-4o": 0.876,                # OpenAI 발표
        "o1": 0.947,                     # OpenAI o1 발표
        "gemini-2.5-pro": 0.895,        # Google 발표
        "med-gemini": 0.91,              # Google Med-Gemini
        "meditron-70b": 0.72,            # EPFL 논문
        "med42-70b": 0.85,               # M42 발표
    },
    "pubmedqa": {
        "claude-opus-4-5": 0.788,
        "gpt-4o": 0.759,
        "o1": 0.82,
        "meditron-70b": 0.82,            # 특화 모델 우세
        "med-gemini": 0.81,
    },
    "mmlu_medical": {
        "gpt-4o": 0.87,
        "claude-opus-4-5": 0.89,
        "gemini-2.5-pro": 0.88,
        "meditron-70b": 0.75,
    },
    # ... 다른 벤치
}


def compute_bench_score(results: list[BenchAnswerResult]) -> dict:
    """정확도 + 95% CI (Wilson score interval)."""
    n = len(results)
    if n == 0:
        return {"accuracy": 0.0, "ci_low": 0.0, "ci_high": 0.0, "n": 0}
    correct = sum(1 for r in results if r.is_correct)
    p = correct / n
    z = 1.96
    denominator = 1 + z**2 / n
    center = (p + z**2 / (2 * n)) / denominator
    margin = z * np.sqrt(p * (1 - p) / n + z**2 / (4 * n**2)) / denominator
    return {
        "accuracy": p,
        "ci_low": max(0, center - margin),
        "ci_high": min(1, center + margin),
        "n": n,
        "correct": correct,
    }


def delta_report(all_results: dict[tuple[str, str], list[BenchAnswerResult]]) -> pd.DataFrame:
    """(bench, model) → 실측 vs 발표 delta."""
    rows = []
    for (bench, model), results in all_results.items():
        score = compute_bench_score(results)
        # 발표 숫자에서 model alias 매칭
        model_key = model.split("-2024")[0].split("-2025")[0]  # snapshot 제거
        published = PUBLISHED_NUMBERS.get(bench, {}).get(model_key)
        delta = (score["accuracy"] - published) if published is not None else None
        
        reproducibility = "N/A"
        if delta is not None:
            if abs(delta) < 0.02:
                reproducibility = "재현됨 (Δ < 2%p)"
            elif abs(delta) < 0.05:
                reproducibility = "부분 재현 (Δ 2~5%p)"
            else:
                reproducibility = f"재현 실패 (Δ {delta:+.1%})"
        
        rows.append({
            "bench": bench,
            "model": model,
            "n": score["n"],
            "correct": score["correct"],
            "accuracy_measured": round(score["accuracy"], 4),
            "ci_95": f"[{score['ci_low']:.3f}, {score['ci_high']:.3f}]",
            "published": published,
            "delta": round(delta, 4) if delta is not None else None,
            "reproducibility": reproducibility,
        })
    return pd.DataFrame(rows)


def reproducibility_heatmap(df: pd.DataFrame, output_path: str = "reproducibility.png") -> None:
    """벤치 × 모델 heatmap (delta 색상)."""
    import matplotlib.pyplot as plt
    import seaborn as sns
    pivot = df.pivot(index="model", columns="bench", values="delta")
    fig, ax = plt.subplots(figsize=(12, 8))
    sns.heatmap(
        pivot, annot=True, cmap="RdYlGn_r", center=0, fmt=".3f",
        cbar_kws={"label": "실측 - 발표 (양수=실측 우수)"},
        ax=ax,
    )
    ax.set_title("Med-LLM 벤치마크 재현 delta 매트릭스\n(양수=실측 우수, 음수=재현 실패)")
    plt.tight_layout()
    plt.savefig(output_path, dpi=150)

def diagnose_reproduction_failure(
    bench: str,
    model: str,
    delta: float,
    all_results: list[BenchAnswerResult],
) -> dict:
    """재현 실패 원인 후보 진단."""
    diagnosis = {
        "bench": bench, "model": model, "delta": delta,
        "possible_causes": [],
    }
    
    if abs(delta) < 0.03:
        diagnosis["possible_causes"].append("재현됨. 원인 조사 불필요.")
        return diagnosis
    
    # 1. Prompt 차이 후보
    prompt_types = {r.metadata.get("prompt_type", "unknown") for r in all_results}
    if "zero_shot_cot" not in prompt_types:
        diagnosis["possible_causes"].append(
            "CoT (Chain-of-Thought) 프롬프트 미사용. 발표 논문은 CoT 사용했을 가능성. zero_shot_cot 프롬프트로 재실행 필요."
        )
    if "few_shot" not in prompt_types:
        diagnosis["possible_causes"].append(
            "Few-shot 예시 미포함. 발표는 5-shot · 25-shot 등 사용했을 가능성."
        )
    
    # 2. 데이터 subset 차이
    total_n = len(all_results)
    diagnosis["possible_causes"].append(
        f"현재 sample n={total_n}. 발표는 full test set (수천 개)에서 측정했을 가능성. sample 확대 필요."
    )
    
    # 3. Self-consistency 미사용
    diagnosis["possible_causes"].append(
        "Self-consistency (majority vote of N samples) 미사용. 발표는 N=5, 10 등 사용했을 가능성."
    )
    
    # 4. Data contamination
    diagnosis["possible_causes"].append(
        f"{bench}는 오픈 데이터. 모델 pre-training corpus에 포함됐을 가능성 (data contamination). "
        "최신 held-out 벤치로 검증 권장."
    )
    
    # 5. 모델 snapshot 차이
    snapshots = {r.model_snapshot for r in all_results}
    diagnosis["possible_causes"].append(
        f"실측 snapshot: {snapshots}. 발표 시점 snapshot과 다를 수 있음."
    )
    
    return diagnosis

async def full_reproduction_bench(
    benches: list[str],       # ["medqa", "pubmedqa", "mmlu", "medmcqa"]
    adapters: list,            # 편 09의 VendorAdapter 리스트
    output_dir: Path,
    sample_size: int | None = 500,
    prompt_types: list[str] = ["zero_shot"],
) -> pd.DataFrame:
    """전체 재현 벤치 실행."""
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # 벤치 로딩
    all_questions = {}
    if "medqa" in benches:
        all_questions["medqa"] = load_medqa(max_samples=sample_size)
    if "pubmedqa" in benches:
        all_questions["pubmedqa"] = load_pubmedqa(max_samples=sample_size)
    if "mmlu" in benches:
        all_questions["mmlu_medical"] = load_mmlu_medical(max_samples=sample_size)
    if "medmcqa" in benches:
        all_questions["medmcqa"] = load_medmcqa(max_samples=sample_size)
    
    print(f"로딩 완료. 벤치별 sample: " +
          ", ".join(f"{k}={len(v)}" for k, v in all_questions.items()))
    
    all_results = {}
    for adapter in adapters:
        for bench_name, questions in all_questions.items():
            for pt in prompt_types:
                print(f"[{adapter.model}] {bench_name} ({pt})")
                results = await run_bench_on_vendor(questions, adapter, prompt_type=pt)
                key = (bench_name, adapter.model)
                if pt != "zero_shot":
                    key = (f"{bench_name}_{pt}", adapter.model)
                all_results[key] = results
    
    # 리포트
    df = delta_report(all_results)
    df.to_csv(output_dir / "reproduction_report.csv", index=False)
    reproducibility_heatmap(df, str(output_dir / "reproducibility_heatmap.png"))
    
    # 재현 실패 진단
    diagnoses = []
    for (bench, model), results in all_results.items():
        row = df[(df["bench"] == bench) & (df["model"] == model)]
        if not row.empty and row.iloc[0]["delta"] is not None:
            diag = diagnose_reproduction_failure(bench, model, row.iloc[0]["delta"], results)
            diagnoses.append(diag)
    
    with open(output_dir / "failure_diagnoses.json", "w", encoding="utf-8") as f:
        import json
        json.dump(diagnoses, f, ensure_ascii=False, indent=2)
    
    return df