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       # lista si es opción múltiple, None si es respuesta abierta
    gold_answer: str                 # para opción múltiple "A"/"B"/..., para respuesta abierta JSON de rubric
    metadata: dict = field(default_factory=dict)


def load_medqa(split: str = "test", max_samples: int | None = None) -> list[BenchQuestion]:
    """Carga de MedQA (estilo USMLE). Repositorio 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"Contexto: {context}\n\nPregunta: {row['question']}\nLa respuesta debe ser una de las siguientes: sí/no/tal vez:",
            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]:
    """Subconjunto médico de MMLU. 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 (examen de licencia médica de la India). 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). Verifique previamente el repositorio y la licencia.

    El nombre real del repositorio debe confirmarse en el momento del anuncio. Aquí se presenta como un esqueleto conceptual.
    """
    try:
        ds = load_dataset("openai/healthbench", split=split)
    except Exception:
        # Alternativa: carga manual de escenarios de muestra del apéndice del artículo
        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 = """Responda la siguiente pregunta del examen nacional de medicina.

Pregunta: {question}

Opciones:
{choices_formatted}

Respuesta (solo una letra: A/B/C/D/E):
"""


ZERO_SHOT_COT_MEDQA = """Responda la siguiente pregunta del examen nacional de medicina. Primero, razone paso a paso y luego proporcione la respuesta final.

Problema: {question}

Opciones:
{choices_formatted}

Razonamiento paso a paso:
"""


FEW_SHOT_MEDQA = """A continuación se presentan ejemplos de preguntas y respuestas del examen nacional de medicina.

Ejemplo 1:
Problema: Paciente masculino de 65 años, con dolor torácico súbito, sudoración y náuseas hace 3 horas. Elevación del segmento ST en el ECG. ¿Cuál es el diagnóstico más probable?
Opciones: A. Infarto agudo de miocardio B. Angina inestable C. Disección aórtica D. Embolia pulmonar E. Pericarditis
Respuesta: A

Ejemplo 2:
Problema: Mujer de 45 años, con sudoración nocturna, pérdida de peso y tos crónica durante 3 meses. Nódulo en el lóbulo superior derecho en la radiografía de tórax. ¿Cuál es el examen prioritario a realizar?
Opciones: A. Tomografía computarizada B. Prueba de baciloscopia en esputo C. Cultivo para micobacterias D. Broncoscopia E. Prueba de Mantoux
Respuesta: B

Ahora responda la siguiente pregunta.

Pregunta: {question}

Opciones:
{choices_formatted}

Respuesta (solo una letra):
"""


ZERO_SHOT_PUBMEDQA = """{question}
"""


ZERO_SHOT_MMLU = """Responda la siguiente pregunta.

Pregunta: {question}

Opciones:
A. {a}
B. {b}
C. {c}
D. {d}

Respuesta (solo una letra de A/B/C/D):
"""


def format_question(q: BenchQuestion, prompt_type: str = "zero_shot") -> str:
    """Renderizado de prompts por benchmark."""
    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  # respuesta abierta

    else:
        raise ValueError(f"Unknown bench: {q.bench_name}")

import re


def extract_choice(response: str, choices: list[str] | None) -> str | None:
    """Extrae una opción A/B/C/D de la respuesta del LLM.

    También maneja respuestas CoT (razonamiento largo seguido de respuesta).
    """
    if not choices:
        return None
    upper = response.upper()
    # Prioritize patterns like "Answer: A" · "The correct answer is A" · "The answer is A"
    for pattern in [
        r"answer[:\s]+([A-E])",
        r"correct answer[is:\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: Last standalone A~E appearing
    matches = re.findall(r"\b([A-E])\b", upper)
    if matches:
        return matches[-1]
    # Fallback: Matching option text
    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:
    """Extract 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

You are an expert in clinical diagnosis and treatment evaluation.
Evaluate the following scenario and response according to the rubric.

Scenario:
{scenario}

Response:
{response}

Rubric (JSON, each item evaluated as 0 or 1):
{rubric}

Output scores for each rubric item in JSON as 0 (not met) · 1 (met):
{{
  "criterion_1": 0 or 1,
  "criterion_2": 0 or 1,
  ...
  "total_score": total,
  "max_score": número total de elementos del rúbrica
}}
"""


def judge_healthbench_response(
    scenario: str,
    response: str,
    rubric_json: str,
    judge_client,
    judge_model: str = "claude-opus-4-5-20250219",
) -> dict:
    """Puntuación del rúbrica con LLM como juez."""
    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": "Error al analizar el juicio"}
    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,  # VendorAdapter del capítulo 09
    prompt_type: str = "zero_shot",
    max_concurrent: int = 5,
) -> list[BenchAnswerResult]:
    """Ejecutar un benchmark × un modelo."""
    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 y otros de respuesta libre requieren juez separado
                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


# Número de vendedores · publicaciones en conferencias (verificar al momento de la ejecución)
PUBLISHED_NUMBERS = {
    "medqa": {
        "claude-opus-4-5": 0.905,      # Ejemplo, anuncio de Anthropic
        "gpt-4o": 0.876,                # Anuncio de OpenAI
        "o1": 0.947,                     # Anuncio de OpenAI o1
        "gemini-2.5-pro": 0.895,        # Anuncio de Google
        "med-gemini": 0.91,              # Google Med-Gemini
        "meditron-70b": 0.72,            # Artículo de EPFL
        "med42-70b": 0.85,               # Presentación M42
    },
    "pubmedqa": {
        "claude-opus-4-5": 0.788,
        "gpt-4o": 0.759,
        "o1": 0.82,
        "meditron-70b": 0.82,            # Dominio de modelos especializados
        "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,
    },
    # ... otros benchmarks
}


def compute_bench_score(results: list[BenchAnswerResult]) -> dict:
    """Precisión + IC del 95% (intervalo de Wilson)."""
    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) → delta entre medido y anunciado."""
    rows = []
    for (bench, model), results in all_results.items():
        score = compute_bench_score(results)
        # Coincidencia de alias del modelo con números anunciados
        model_key = model.split("-2024")[0].split("-2025")[0]  # Eliminar instantánea
        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 = "Reproducido (Δ < 2%p)"
            elif abs(delta) < 0.05:
                reproducibility = "Parcialmente reproducido (Δ 2~5%p)"
            else:
                reproducibility = f"Fallo de reproducción (Δ {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:
    """Mapa de calor benchmark × modelo (color 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": "Medición - Publicado (positivo=mejor medición)"},
        ax=ax,
    )
    ax.set_title("Métrica delta de reproducción del benchmark Med-LLM\n(positivo=mejor medición, negativo=fallo de reproducció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:
    """Diagnóstico de posibles causas de fallo en la reproducción."""
    diagnosis = {
        "bench": bench, "model": model, "delta": delta,
        "possible_causes": [],
    }

    if abs(delta) < 0.03:
        diagnosis["possible_causes"].append("Reproducido. No es necesaria una investigación de las causas.")
        return diagnosis

    # 1. Posible diferencia en el 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(
            "No se utilizó el prompt CoT (Chain-of-Thought). Es posible que el artículo publicado haya utilizado CoT. Necesario reejecutar con el prompt zero_shot_cot."
        )
    if "few_shot" not in prompt_types:
        diagnosis["possible_causes"].append(
            "Sin ejemplos few-shot incluidos. Es posible que la publicación haya utilizado 5-shot, 25-shot, etc."
        )

    # 2. Diferencia en el subconjunto de datos
    total_n = len(all_results)
    diagnosis["possible_causes"].append(
        f"n actual de muestras={total_n}. Es posible que la publicación se haya medido en el conjunto completo de pruebas (miles). Necesario ampliar las muestras."
    )

    # 3. No se utilizó la consistencia interna
    diagnosis["possible_causes"].append(
        "No se utilizó la consistencia interna (voto mayoritario de N muestras). Es posible que la publicación haya utilizado N=5, 10, etc."
    )

    # 4. Data contamination
    diagnosis["possible_causes"].append(
        f"{bench} es un conjunto de datos abierto. Posible inclusión en el corpus de preentrenamiento del modelo (contaminación de datos). "
        "Se recomienda validar con benchmarks recientes no utilizados en el entrenamiento."
    )

    # 5. Diferencia en las instantáneas del modelo
    snapshots = {r.model_snapshot for r in all_results}
    diagnosis["possible_causes"].append(
        f"Instantánea observada: {snapshots}. Puede diferir de la instantánea publicada en el momento del anuncio."
    )

    return diagnosis

async def full_reproduction_bench(
    benches: list[str],       # ["medqa", "pubmedqa", "mmlu", "medmcqa"]
    adapters: list,            # Lista de VendorAdapter del capítulo 09
    output_dir: Path,
    sample_size: int | None = 500,
    prompt_types: list[str] = ["zero_shot"],
) -> pd.DataFrame:
    """Ejecución completa del benchmark de reproducción."""
    output_dir.mkdir(parents=True, exist_ok=True)

    # Carga de benchmarks
    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"Carga completada. Muestras por benchmark: " +
          ", ".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

    # Informe
    df = delta_report(all_results)
    df.to_csv(output_dir / "reproduction_report.csv", index=False)
    reproducibility_heatmap(df, str(output_dir / "reproducibility_heatmap.png"))

    # Diagnóstico de fallo en la reproducción
    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