"""
Benchmarking Frontier LMMs for PNI Extraction from GI Pathology Reports
=======================================================================
This script implements the complete study pipeline described in:
"Comparative Performance of Frontier Large Language Models for Extracting
High-Risk Pathologic Features from Unstructured Gastrointestinal Oncology Reports"

Requirements:
    pip install openai google-generativeai anthropic transformers torch
             scikit-learn scipy pandas numpy nltk rouge-score sacrebleu
             statsmodels matplotlib seaborn tqdm
"""
from dotenv import load_dotenv
load_dotenv()


import os
import json
import time
import re
import pandas as pd
import numpy as np
from typing import Dict, List, Tuple, Optional, Any
from dataclasses import dataclass, field
from collections import defaultdict
from tqdm import tqdm

# API clients
import openai
import google.generativeai as genai
import ollama
import anthropic
import requests

# ML and metrics
from sklearn.model_selection import train_test_split
from sklearn.metrics import (
    accuracy_score, precision_score, recall_score, f1_score,
    roc_auc_score, roc_curve, confusion_matrix
)
from scipy import stats
from statsmodels.stats.multitest import multipletests
import torch
from transformers import (
    AutoTokenizer, AutoModelForSequenceClassification,
    Trainer, TrainingArguments, DataCollatorWithPadding
)
from datasets import Dataset

# NLP metrics
from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction
from rouge_score import rouge_scorer

# Visualization
import matplotlib.pyplot as plt
import seaborn as sns

# =============================================================================
# CONFIGURATION (Replace with your actual values)
# =============================================================================

@dataclass
class Config:
    # Data paths
    data_path: str = "D:\\BattleofBigs\\TCGA_Reports_PNI_Labeled.csv"
    output_dir: str = "./benchmark_results"
    
    # API Keys (set via environment variables or directly)
    # openai_api_key: str = os.getenv("OPENAI_API_KEY", "")
    # gemini_api_key: str = os.getenv("GEMINI_API_KEY", "")
    # anthropic_api_key: str = os.getenv("ANTHROPIC_API_KEY", "")
    deepseek_api_key: str = os.getenv("DEEPSEEK_API_KEY", "")
    
    # Model versions (exact API identifiers)
    # gpt4o_model: str = "gpt-4o-2024-08-06"
    # gemini_model: str = "gemini-2.5-pro-exp-03-25"
    # claude_model: str = "claude-3-7-sonnet-20250219"
    deepseek_model: str = "deepseek-reasoner"
    
    # Experiment parameters
    random_seed: int = 42
    test_size: float = 0.2
    val_size: float = 0.1
    max_tokens: int = 500
    temperature: float = 0.0
    human_time_limit_seconds: int = 90
    bootstrap_iterations: int = 10000
    confidence_level: float = 0.95
    
    # BioBERT training
    biobert_model_name: str = "dmis-lab/biobert-base-cased-v1.2"
    batch_size: int = 16
    learning_rate: float = 2e-5
    num_epochs: int = 3

config = Config()

# =============================================================================
# DATA LOADING AND PREPROCESSING
# =============================================================================

def load_data(filepath: str) -> pd.DataFrame:
    """
    Load annotated pathology reports.
    Expected columns:
        - report_id: unique identifier
        - report_text: full narrative pathology report
        - pni_ground_truth: 'Present', 'Absent', or 'Indeterminate'
        - cancer_type: e.g., 'Colorectal', 'Gastric', 'Pancreatic'
        - gold_standard_uncertainty: boolean flag (True if ambiguous)
        - resident_annotation_1, resident_annotation_2, resident_annotation_3: human baseline
    """
    df = pd.read_csv(filepath)
    required_cols = ['report_id', 'report_text', 'pni_ground_truth']
    for col in required_cols:
        assert col in df.columns, f"Missing required column: {col}"
    
    # Map ground truth to binary for AUC: Present=1, Absent/Indeterminate=0
    df['pni_binary'] = df['pni_ground_truth'].apply(
        lambda x: 1 if x == 'Present' else 0
    )
    return df

def split_data(df: pd.DataFrame, test_size: float, val_size: float, seed: int):
    """Stratified split maintaining class distribution."""
    # First split: train+val vs test
    train_val_df, test_df = train_test_split(
        df, test_size=test_size, stratify=df['pni_ground_truth'],
        random_state=seed
    )
    # Second split: train vs val
    val_ratio = val_size / (1 - test_size)
    train_df, val_df = train_test_split(
        train_val_df, test_size=val_ratio,
        stratify=train_val_df['pni_ground_truth'],
        random_state=seed
    )
    
    return train_df, val_df, test_df

# =============================================================================
# PROMPT ENGINEERING (Two-Stage Chain-of-Thought + JSON)
# =============================================================================

SYSTEM_PROMPT = """You are an expert gastrointestinal pathologist assistant. 
Your task is to extract perineural invasion (PNI) status from pathology reports 
with high precision. Follow the instructions exactly."""

def build_prompt(report_text: str) -> str:
    """Construct the two-stage prompt."""
    return f"""Analyze the following pathology report and determine the perineural invasion (PNI) status.

Step 1 - Reasoning (think step by step):
First, identify all mentions of perineural invasion, neural involvement, or nerve involvement anywhere in the report—including the microscopic description, not only the final diagnosis. Look for synonyms like "perineural," "neural invasion," "nerve involvement," "intraneural," "extraneural."
Second, determine whether any qualifying negation (e.g., 'no evidence of,' 'not identified,' 'negative for,' 'without') applies.
Third, classify PNI status as: Present, Absent, or Indeterminate.

Step 2 - Structured Output:
Provide your final answer in the following JSON format only, with no additional text:
{{
  "PNI_status": "Present|Absent|Indeterminate",
  "evidence_quote": "exact phrase from report supporting your decision",
  "confidence": "High|Moderate|Low"
}}

Pathology Report:
---
{report_text}
---"""

# =============================================================================
# LMM API CLIENTS
# =============================================================================

class LMMClient:
    """Base class for LMM inference."""
    def __init__(self, model_name: str):
        self.model_name = model_name
    
    def query(self, prompt: str) -> Tuple[str, float]:
        """Return (response_text, latency_seconds)."""
        raise NotImplementedError

class GPT4oClient(LMMClient):
    def __init__(self, api_key: str, model: str):
        super().__init__(model)
        self.client = openai.OpenAI(api_key=api_key)
    
    def query(self, prompt: str) -> Tuple[str, float]:
        start = time.time()
        response = self.client.chat.completions.create(
            model=self.model_name,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": prompt}
            ],
            temperature=config.temperature,
            max_tokens=config.max_tokens
        )
        latency = time.time() - start
        return response.choices[0].message.content, latency
    
class OllamaClient(LMMClient):
    def __init__(self, model_name: str):
        super().__init__(model_name)
        # No API key needed!
    
    def query(self, prompt: str) -> Tuple[str, float]:
        start = time.time()
        
        # Construct the messages in the format Ollama expects
        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": prompt}
        ]
        
        # Send the request to the local Ollama instance
        response = ollama.chat(
            model=self.model_name, # This will be 'llama3.3'
            messages=messages,
            options={
                'temperature': config.temperature,
                'num_predict': config.max_tokens,
            }
        )
        
        latency = time.time() - start
        # Extract the text content from the response
        return response['message']['content'], latency
    
class GeminiClient(LMMClient):
    def __init__(self, api_key: str, model: str):
        super().__init__(model)
        genai.configure(api_key=api_key)
        self.model = genai.GenerativeModel(model)
    
    def query(self, prompt: str) -> Tuple[str, float]:
        start = time.time()
        response = self.model.generate_content(
            [SYSTEM_PROMPT, prompt],
            generation_config=genai.types.GenerationConfig(
                temperature=config.temperature,
                max_output_tokens=config.max_tokens
            )
        )
        latency = time.time() - start
        return response.text, latency

class ClaudeClient(LMMClient):
    def __init__(self, api_key: str, model: str):
        super().__init__(model)
        self.client = anthropic.Anthropic(api_key=api_key)
    
    def query(self, prompt: str) -> Tuple[str, float]:
        start = time.time()
        response = self.client.messages.create(
            model=self.model_name,
            system=SYSTEM_PROMPT,
            messages=[{"role": "user", "content": prompt}],
            temperature=config.temperature,
            max_tokens=config.max_tokens
        )
        latency = time.time() - start
        return response.content[0].text, latency

class DeepSeekClient(LMMClient):
    def __init__(self, api_key: str, model: str):
        super().__init__(model)
        self.api_key = api_key
        self.base_url = "https://api.deepseek.com/v1/chat/completions"
    
    def query(self, prompt: str) -> Tuple[str, float]:
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        payload = {
            "model": self.model_name,
            "messages": [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": prompt}
            ],
            "temperature": config.temperature,
            "max_tokens": config.max_tokens
        }
        start = time.time()
        response = requests.post(self.base_url, headers=headers, json=payload)
        latency = time.time() - start
        response.raise_for_status()
        return response.json()["choices"][0]["message"]["content"], latency

# =============================================================================
# RESPONSE PARSING AND VALIDATION
# =============================================================================

def extract_json_from_response(response: str) -> Optional[Dict]:
    """Extract JSON object from model response, handling markdown code blocks."""
    # Try to find JSON in code blocks
    json_pattern = r'```(?:json)?\s*(\{.*?\})\s*```'
    match = re.search(json_pattern, response, re.DOTALL)
    if match:
        json_str = match.group(1)
    else:
        # Try to find raw JSON
        json_match = re.search(r'\{.*\}', response, re.DOTALL)
        if json_match:
            json_str = json_match.group(0)
        else:
            return None
    
    try:
        return json.loads(json_str)
    except json.JSONDecodeError:
        return None

def parse_pni_prediction(response: str) -> Dict[str, Any]:
    """Extract PNI status, evidence, confidence, and detect hallucinations."""
    parsed = extract_json_from_response(response)
    result = {
        "raw_response": response,
        "PNI_status": None,
        "evidence_quote": None,
        "confidence": None,
        "parse_error": True,
        "hallucination_flag": False
    }
    
    if parsed:
        result["parse_error"] = False
        result["PNI_status"] = parsed.get("PNI_status")
        result["evidence_quote"] = parsed.get("evidence_quote")
        result["confidence"] = parsed.get("confidence")
        
        # Normalize status
        if result["PNI_status"]:
            status_lower = result["PNI_status"].lower()
            if "present" in status_lower:
                result["PNI_status"] = "Present"
            elif "absent" in status_lower:
                result["PNI_status"] = "Absent"
            elif "indeterminate" in status_lower:
                result["PNI_status"] = "Indeterminate"
    
    return result

# =============================================================================
# FAILURE MODE CLASSIFICATION
# =============================================================================

class FailureModeClassifier:
    """Classify errors into taxonomy categories."""
    
    NEGATION_PATTERNS = [
        r'no\s+evidence\s+of', r'not\s+identified', r'negative\s+for',
        r'without', r'absent', r'free\s+of', r'no\s+definite',
        r'does\s+not\s+show', r'failed\s+to\s+demonstrate'
    ]
    HEDGING_PATTERNS = [
        r'possible', r'suspected', r'suspicious', r'likely',
        r'probable', r'cannot\s+exclude', r'suggestive', r'equivocal'
    ]
    SYNONYM_MAP = {
        'scirrhous': 'ductal',
        'neural': 'perineural',
        'intraneural': 'perineural',
        'extraneural': 'perineural'
    }
    
    @classmethod
    def classify(cls, report_text: str, prediction: str, ground_truth: str,
                 evidence_quote: Optional[str] = None) -> str:
        """
        Returns failure mode category or None if correct.
        """
        if prediction == ground_truth:
            return None
        
        report_lower = report_text.lower()
        
        # Negation failure: model missed negation
        has_negation = any(re.search(pat, report_lower) for pat in cls.NEGATION_PATTERNS)
        if has_negation and prediction == "Present" and ground_truth == "Absent":
            return "Negation Failure"
        
        # Hedging ambiguity: model couldn't commit on probabilistic language
        has_hedging = any(re.search(pat, report_lower) for pat in cls.HEDGING_PATTERNS)
        if has_hedging and prediction in ["Indeterminate", None]:
            return "Hedging Ambiguity"
        
        # Synonym blindness: missed uncommon term
        if evidence_quote:
            quote_lower = evidence_quote.lower()
            for syn, std in cls.SYNONYM_MAP.items():
                if syn in quote_lower and std not in quote_lower:
                    return "Synonym Blindness"
        
        # Anchoring bias: check if later correction missed (simplified)
        if "addendum" in report_lower and "corrected" in report_lower:
            # If model gave wrong answer and addendum exists, possible anchoring
            return "Anchoring Bias"
        
        return "Other"

# =============================================================================
# METRICS COMPUTATION
# =============================================================================

def bootstrap_ci(data: np.ndarray, func: callable, n_iter: int = 10000,
                 alpha: float = 0.05) -> Tuple[float, float, float]:
    """Compute bootstrap confidence interval for a metric."""
    np.random.seed(config.random_seed)
    estimates = []
    n = len(data)
    for _ in range(n_iter):
        sample = np.random.choice(data, size=n, replace=True)
        estimates.append(func(sample))
    ci_lower = np.percentile(estimates, 100 * alpha / 2)
    ci_upper = np.percentile(estimates, 100 * (1 - alpha / 2))
    point_estimate = func(data)
    return point_estimate, ci_lower, ci_upper

def compute_metrics(y_true: List[str], y_pred: List[str],
                    y_scores: Optional[List[float]] = None,
                    reference_texts: Optional[List[str]] = None,
                    generated_texts: Optional[List[str]] = None) -> Dict:
    """Compute all primary and secondary metrics."""
    # Binary conversion for AUC/sensitivity/specificity
    y_true_bin = [1 if x == "Present" else 0 for x in y_true]
    y_pred_bin = [1 if x == "Present" else 0 for x in y_pred]
    
    metrics = {
        "accuracy": accuracy_score(y_true, y_pred),
        "sensitivity": recall_score(y_true_bin, y_pred_bin, pos_label=1),
        "specificity": recall_score(y_true_bin, y_pred_bin, pos_label=0),
        "f1": f1_score(y_true_bin, y_pred_bin, zero_division=0)
    }
    
    if y_scores:
        metrics["auc_roc"] = roc_auc_score(y_true_bin, y_scores)
    
    # BLEU and ROUGE
    if reference_texts and generated_texts:
        smoother = SmoothingFunction().method1
        bleu_scores = []
        for ref, gen in zip(reference_texts, generated_texts):
            if ref and gen:
                bleu_scores.append(sentence_bleu([ref.split()], gen.split(),
                                                 smoothing_function=smoother))
        metrics["bleu"] = np.mean(bleu_scores) if bleu_scores else 0.0
        
        scorer = rouge_scorer.RougeScorer(['rougeL'], use_stemmer=True)
        rouge_scores = [scorer.score(ref, gen)['rougeL'].fmeasure
                        for ref, gen in zip(reference_texts, generated_texts) if ref and gen]
        metrics["rouge_l"] = np.mean(rouge_scores) if rouge_scores else 0.0
    
    return metrics

def mcnemar_test(y_true: List[str], y_pred1: List[str], y_pred2: List[str]) -> Dict:
    """McNemar's test for paired nominal data."""
    # Create contingency table
    both_correct = sum(1 for t, p1, p2 in zip(y_true, y_pred1, y_pred2)
                       if p1 == t and p2 == t)
    only1_correct = sum(1 for t, p1, p2 in zip(y_true, y_pred1, y_pred2)
                        if p1 == t and p2 != t)
    only2_correct = sum(1 for t, p1, p2 in zip(y_true, y_pred1, y_pred2)
                        if p1 != t and p2 == t)
    both_wrong = sum(1 for t, p1, p2 in zip(y_true, y_pred1, y_pred2)
                     if p1 != t and p2 != t)
    
    table = [[both_correct, only1_correct],
             [only2_correct, both_wrong]]
    
    # McNemar's test uses the discordant pairs
    if only1_correct + only2_correct > 0:
        stat, p_value = stats.binomtest(min(only1_correct, only2_correct),
                                        n=only1_correct + only2_correct, p=0.5).pvalue, None
        # Actually use exact binomial test
        from scipy.stats import binomtest
        result = binomtest(min(only1_correct, only2_correct),
                           n=only1_correct + only2_correct, p=0.5)
        p_value = result.pvalue
    else:
        stat, p_value = 0, 1.0
    
    return {
        "statistic": stat,
        "p_value": p_value,
        "discordant_pairs": only1_correct + only2_correct,
        "delta_accuracy": (only1_correct - only2_correct) / len(y_true)
    }

def delong_test(y_true: np.ndarray, y_score1: np.ndarray, y_score2: np.ndarray) -> float:
    """DeLong's test for comparing two AUCs (simplified implementation)."""
    # This is a placeholder - for actual use, consider using the 'roc' package or
    # implementing the full DeLong method. Here we approximate with bootstrap.
    from sklearn.utils import resample
    
    np.random.seed(config.random_seed)
    auc_diff = []
    n = len(y_true)
    for _ in range(1000):
        idx = resample(np.arange(n), n_samples=n)
        auc1 = roc_auc_score(y_true[idx], y_score1[idx])
        auc2 = roc_auc_score(y_true[idx], y_score2[idx])
        auc_diff.append(auc1 - auc2)
    
    # Two-tailed p-value from bootstrap distribution
    p_value = 2 * min(np.mean(np.array(auc_diff) <= 0), np.mean(np.array(auc_diff) >= 0))
    return p_value

# =============================================================================
# BIOBERT BASELINE IMPLEMENTATION
# =============================================================================

def train_biobert(train_df: pd.DataFrame, val_df: pd.DataFrame) -> Any:
    """Fine-tune BioBERT for PNI classification."""
    tokenizer = AutoTokenizer.from_pretrained(config.biobert_model_name)
    model = AutoModelForSequenceClassification.from_pretrained(
        config.biobert_model_name, num_labels=2
    )
    
    def tokenize_function(examples):
        return tokenizer(examples["report_text"], truncation=True, padding=False, max_length=512)
    
    # Convert to HuggingFace Dataset
    train_dataset = Dataset.from_pandas(train_df[['report_text', 'pni_binary']])
    val_dataset = Dataset.from_pandas(val_df[['report_text', 'pni_binary']])
    
    train_dataset = train_dataset.map(tokenize_function, batched=True)
    val_dataset = val_dataset.map(tokenize_function, batched=True)
    
    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
    
    training_args = TrainingArguments(
        output_dir=os.path.join(config.output_dir, "biobert_checkpoints"),
        evaluation_strategy="epoch",
        save_strategy="epoch",
        learning_rate=config.learning_rate,
        per_device_train_batch_size=config.batch_size,
        per_device_eval_batch_size=config.batch_size,
        num_train_epochs=config.num_epochs,
        weight_decay=0.01,
        load_best_model_at_end=True,
        metric_for_best_model="accuracy",
        seed=config.random_seed,
        logging_dir=os.path.join(config.output_dir, "logs"),
        report_to="none"
    )
    
    def compute_metrics(eval_pred):
        logits, labels = eval_pred
        predictions = np.argmax(logits, axis=-1)
        return {"accuracy": accuracy_score(labels, predictions)}
    
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=val_dataset,
        data_collator=data_collator,
        compute_metrics=compute_metrics,
        tokenizer=tokenizer
    )
    
    trainer.train()
    return trainer, tokenizer

def evaluate_biobert(trainer: Trainer, tokenizer, test_df: pd.DataFrame) -> Tuple[List[str], List[float]]:
    """Generate predictions and confidence scores from fine-tuned BioBERT."""
    test_texts = test_df['report_text'].tolist()
    inputs = tokenizer(test_texts, truncation=True, padding=True, return_tensors="pt", max_length=512)
    
    trainer.model.eval()
    with torch.no_grad():
        outputs = trainer.model(**inputs)
        logits = outputs.logits
        probs = torch.softmax(logits, dim=-1)
        preds = torch.argmax(logits, dim=-1).numpy()
        scores = probs[:, 1].numpy()  # probability of class 1 (Present)
    
    pred_labels = ["Present" if p == 1 else "Absent" for p in preds]
    return pred_labels, scores.tolist()

# =============================================================================
# MAIN EXECUTION PIPELINE
# =============================================================================

def run_benchmark():
    """Execute the full benchmarking study."""
    os.makedirs(config.output_dir, exist_ok=True)
    
    # 1. Load and split data
    print("Loading data...")
    df = load_data(config.data_path)
    train_df, val_df, test_df = split_data(df, config.test_size, config.val_size, config.random_seed)
    print(f"Data split: Train={len(train_df)}, Val={len(val_df)}, Test={len(test_df)}")
    
    # Save dataset characteristics (Table 2)
    table2_data = {
        "Characteristic": ["Total Cases", "Malignant", "Benign", "Cancer Type Distribution",
                          "Gold Standard Uncertainty Cases"],
        "Training Set": [len(train_df), 
                         len(train_df[train_df['pni_ground_truth'] == 'Present']),
                         len(train_df[train_df['pni_ground_truth'] == 'Absent']),
                         str(train_df['cancer_type'].value_counts().to_dict()),
                         len(train_df[train_df.get('gold_standard_uncertainty', False)])],
        "Validation Set": [len(val_df), 
                           len(val_df[val_df['pni_ground_truth'] == 'Present']),
                           len(val_df[val_df['pni_ground_truth'] == 'Absent']),
                           str(val_df['cancer_type'].value_counts().to_dict()),
                           len(val_df[val_df.get('gold_standard_uncertainty', False)])],
        "Test Set": [len(test_df), 
                     len(test_df[test_df['pni_ground_truth'] == 'Present']),
                     len(test_df[test_df['pni_ground_truth'] == 'Absent']),
                     str(test_df['cancer_type'].value_counts().to_dict()),
                     len(test_df[test_df.get('gold_standard_uncertainty', False)])]
    }
    pd.DataFrame(table2_data).to_csv(os.path.join(config.output_dir, "Table2_Dataset_Characteristics.csv"), index=False)
    
    # 2. Initialize LMM clients
    # In your run_benchmark() function
    print("Initializing LMM clients...")
    clients = {
        # "GPT-4o": GPT4oClient(config.openai_api_key, config.gpt4o_model),
        # "Gemini 2.5 Pro": GeminiClient(config.gemini_api_key, config.gemini_model),
        # "Claude 3.7 Sonnet": ClaudeClient(config.anthropic_api_key, config.claude_model),
        "DeepSeek-R1": DeepSeekClient(config.deepseek_api_key, config.deepseek_model),
        # "Llama 3.3 (Local)": OllamaClient("llama3.3") # <-- ADD THIS LINE
    }
    
    # 3. Run LMM inference on test set
    results = {model_name: [] for model_name in clients.keys()}
    latencies = {model_name: [] for model_name in clients.keys()}
    failure_modes = {model_name: [] for model_name in clients.keys()}
    
    for idx, row in tqdm(test_df.iterrows(), total=len(test_df), desc="LMM Inference"):
        prompt = build_prompt(row['report_text'])
        ground_truth = row['pni_ground_truth']
        
        for model_name, client in clients.items():
            response, latency = client.query(prompt)
            parsed = parse_pni_prediction(response)
            
            pred_status = parsed['PNI_status'] if parsed['PNI_status'] else "Indeterminate"
            results[model_name].append({
                "report_id": row['report_id'],
                "ground_truth": ground_truth,
                "prediction": pred_status,
                "raw_response": response,
                "evidence_quote": parsed['evidence_quote'],
                "confidence": parsed['confidence'],
                "parse_error": parsed['parse_error']
            })
            latencies[model_name].append(latency)
            
            # Classify failure mode if incorrect
            fm = FailureModeClassifier.classify(
                row['report_text'], pred_status, ground_truth,
                parsed['evidence_quote']
            )
            failure_modes[model_name].append(fm)
    
    # 4. Human baseline (using provided resident annotations)
    human_results = []
    for idx, row in test_df.iterrows():
        # Majority vote among residents or use provided single annotation
        resident_cols = [c for c in test_df.columns if 'resident_annotation' in c]
        if resident_cols:
            votes = row[resident_cols].tolist()
            pred = max(set(votes), key=votes.count)  # majority vote
        else:
            pred = row.get('resident_consensus', 'Indeterminate')
        human_results.append(pred)
    
    # 5. BioBERT baseline
    print("Training BioBERT baseline...")
    trainer, tokenizer = train_biobert(train_df, val_df)
    biobert_preds, biobert_scores = evaluate_biobert(trainer, tokenizer, test_df)
    
    # 6. Compile predictions for metrics
    test_true = test_df['pni_ground_truth'].tolist()
    test_true_bin = test_df['pni_binary'].tolist()
    
    all_predictions = {}
    all_scores = {}
    
    for model_name in clients.keys():
        preds = [r['prediction'] for r in results[model_name]]
        all_predictions[model_name] = preds
        # For AUC: map confidence to score or use binary prediction
        scores = [1 if p == "Present" else 0 for p in preds]
        all_scores[model_name] = scores
    
    all_predictions["Time-Pressured Residents"] = human_results
    all_predictions["Fine-tuned BioBERT"] = biobert_preds
    all_scores["Fine-tuned BioBERT"] = biobert_scores
    # For residents, no continuous scores
    all_scores["Time-Pressured Residents"] = [1 if p == "Present" else 0 for p in human_results]
    
    # 7. Compute metrics and bootstrap CIs
    metrics_summary = {}
    for model_name, preds in all_predictions.items():
        # For generation metrics, use evidence_quote or raw response as generated text
        gen_texts = [r.get('evidence_quote', '') for r in results.get(model_name, [])] if model_name in clients else []
        ref_texts = [f"PNI {gt}" for gt in test_true]  # simple reference
        
        metrics = compute_metrics(
            test_true, preds,
            y_scores=all_scores.get(model_name),
            reference_texts=ref_texts if gen_texts else None,
            generated_texts=gen_texts if gen_texts else None
        )
        
        # Add bootstrap CIs for key metrics
        if len(set(test_true)) > 1:
            acc_data = np.array([1 if p == t else 0 for p, t in zip(preds, test_true)])
            metrics['accuracy_ci'] = bootstrap_ci(acc_data, np.mean)
            
            if metrics.get('auc_roc'):
                # For AUC, bootstrap the difference (simplified)
                pass
        
        metrics['mean_latency'] = np.mean(latencies.get(model_name, [])) if model_name in clients else None
        metrics['hallucination_rate'] = np.mean([1 for r in results.get(model_name, []) if r.get('parse_error', False)]) if model_name in clients else None
        
        metrics_summary[model_name] = metrics
    
    # 8. Create Table 1 (Comparative Performance)
    table1_rows = []
    for model_name, metrics in metrics_summary.items():
        row = {
            "Model": model_name,
            "Accuracy (%)": f"{metrics['accuracy']*100:.1f}",
            "Sensitivity (%)": f"{metrics['sensitivity']*100:.1f}",
            "Specificity (%)": f"{metrics['specificity']*100:.1f}",
            "AUC-ROC": f"{metrics.get('auc_roc', 0):.2f}",
            "BLEU Score": f"{metrics.get('bleu', 0):.2f}",
            "ROUGE-L": f"{metrics.get('rouge_l', 0):.2f}"
        }
        table1_rows.append(row)
    pd.DataFrame(table1_rows).to_csv(os.path.join(config.output_dir, "Table1_Performance.csv"), index=False)
    
    # 9. Pairwise statistical comparisons (Table 3)
    table3_rows = []
    model_names = list(clients.keys()) + ["Time-Pressured Residents", "Fine-tuned BioBERT"]
    p_values = []
    
    for i, m1 in enumerate(model_names):
        for j, m2 in enumerate(model_names):
            if i >= j:
                continue
            pred1 = all_predictions[m1]
            pred2 = all_predictions[m2]
            # McNemar test for accuracy
            result = mcnemar_test(test_true, pred1, pred2)
            delta_acc = (accuracy_score(test_true, pred1) - accuracy_score(test_true, pred2)) * 100
            p_values.append(result['p_value'])
            table3_rows.append({
                "Model Pair": f"{m1} vs {m2}",
                "Metric": "Accuracy",
                "Δ Difference": f"{delta_acc:.1f}%",
                "95% CI": "—",  # Would compute with bootstrap
                "p-value": f"{result['p_value']:.4f}"
            })
    
    # Apply FDR correction
    _, p_corrected, _, _ = multipletests(p_values, method='fdr_bh')
    for row, p_corr in zip(table3_rows, p_corrected):
        row['p-value (FDR corrected)'] = f"{p_corr:.4f}"
    
    pd.DataFrame(table3_rows).to_csv(os.path.join(config.output_dir, "Table3_Pairwise_Comparisons.csv"), index=False)
    
    # 10. Failure Mode Taxonomy (Table 4)
    failure_counts = {model: defaultdict(int) for model in clients.keys()}
    for model_name in clients.keys():
        for fm in failure_modes[model_name]:
            if fm:
                failure_counts[model_name][fm] += 1
    
    table4_rows = []
    all_failure_types = set()
    for counts in failure_counts.values():
        all_failure_types.update(counts.keys())
    
    for fm_type in sorted(all_failure_types):
        row = {"Failure Type": fm_type}
        for model_name in clients.keys():
            count = failure_counts[model_name][fm_type]
            pct = (count / len(test_df)) * 100
            row[model_name] = f"{pct:.1f}% ({count})"
        table4_rows.append(row)
    
    pd.DataFrame(table4_rows).to_csv(os.path.join(config.output_dir, "Table4_Failure_Modes.csv"), index=False)
    
    # 11. Save all results for reproducibility
    with open(os.path.join(config.output_dir, "full_results.json"), "w") as f:
        json.dump({
            "metrics": metrics_summary,
            "predictions": all_predictions,
            "config": {k: str(v) for k, v in config.__dict__.items()}
        }, f, indent=2)
    
    print(f"\nBenchmark complete. Results saved to {config.output_dir}/")
    print("\nKey Findings:")
    best_model = max(metrics_summary.items(), key=lambda x: x[1]['accuracy'])
    print(f"Best model: {best_model[0]} (Accuracy: {best_model[1]['accuracy']*100:.1f}%)")

if __name__ == "__main__":
    run_benchmark()