evaluator.py

#####################################################################################################################################################################
# import re
# import json
# import torch
# import pandas as pd
# import matplotlib.pyplot as plt
# import seaborn as sns
# import os
# import uuid
# from transformers import AutoTokenizer, AutoModelForSequenceClassification
# from sentence_transformers import SentenceTransformer, util

# import matplotlib.pyplot as plt
# import numpy as np

# def plot_radar_chart(metrics_df, agents, metrics, out_path="/tmp/radar.png"):
#     """
#     Radar chart comparing multiple agents across metrics.
#     """
#     labels = metrics
#     num_vars = len(labels)

#     # Compute angle for each axis
#     angles = np.linspace(0, 2 * np.pi, num_vars, endpoint=False).tolist()
#     angles += angles[:1]  # close loop

#     fig, ax = plt.subplots(figsize=(6, 6), subplot_kw=dict(polar=True))

#     for agent in agents:
#         values = []
#         for m in metrics:
#             mean_val = metrics_df.loc[metrics_df['agent'] == agent, m].mean()
#             values.append(mean_val if not np.isnan(mean_val) else 0)
#         values += values[:1]
#         ax.plot(angles, values, label=agent, linewidth=2)
#         ax.fill(angles, values, alpha=0.25)

#     ax.set_xticks(angles[:-1])
#     ax.set_xticklabels(labels)
#     ax.set_yticklabels([])
#     ax.legend(loc="upper right", bbox_to_anchor=(1.3, 1.1))
#     ax.set_title("Agent Performance Radar Chart")
    
#     plt.tight_layout()
#     plt.savefig(out_path)
#     plt.close()
#     return out_path
    
# import seaborn as sns

# def plot_heatmap(metrics_df, out_path="/tmp/heatmap.png"):
#     pivot = metrics_df.groupby("agent")[
#         ["accuracy", "hallucination", "instruction_following", "coherence", "assumption"]
#     ].mean()
    
#     plt.figure(figsize=(8, 5))
#     sns.heatmap(pivot, annot=True, cmap="viridis", fmt=".2f")
#     plt.title("Agent × Metric Heatmap")
#     plt.tight_layout()
#     plt.savefig(out_path)
#     plt.close()
#     return out_path

# # --------------------------
# # MODEL LOADING
# # --------------------------
# NLI_MODEL = "textattack/roberta-base-MNLI"
# EMBED_MODEL = "sentence-transformers/all-MiniLM-L6-v2"

# # Load NLI model & tokenizer
# nli_tokenizer = AutoTokenizer.from_pretrained(NLI_MODEL)
# nli_model = AutoModelForSequenceClassification.from_pretrained(NLI_MODEL)
# nli_model.to("cpu")
# nli_model.eval()

# # Load embedding model
# embed_model = SentenceTransformer(EMBED_MODEL)

# # Label mapping from config
# id2label = {int(k): v.upper() for k, v in nli_model.config.id2label.items()}


# # --------------------------
# # METRIC FUNCTIONS
# # --------------------------
# def check_instruction_following(prompt: str, response: str) -> float:
#     """Embedding-based similarity between prompt and response."""
#     if not prompt or not response:
#         return 0.0
#     p_emb = embed_model.encode(prompt, convert_to_tensor=True)
#     r_emb = embed_model.encode(response, convert_to_tensor=True)
#     sim = float(util.cos_sim(p_emb, r_emb).item())
#     return round(max(0.0, min(1.0, sim)), 3)


# def check_hallucination(reference: str, response: str) -> float:
#     """
#     Single hallucination score:
#     Entailment prob - Contradiction prob (normalized to [0,1]).
#     Higher = less hallucination.
#     """
#     if not reference or not response:
#         return 0.0
#     with torch.no_grad():
#         inputs = nli_tokenizer.encode_plus(reference, response, return_tensors="pt", truncation=True)
#         outputs = nli_model(**inputs)
#         probs = torch.softmax(outputs.logits, dim=-1).cpu().numpy()[0]

#     entail_prob, contra_prob = 0.0, 0.0
#     for idx, p in enumerate(probs):
#         label = id2label.get(idx, "")
#         if "ENTAIL" in label:
#             entail_prob = float(p)
#         elif "CONTRA" in label:
#             contra_prob = float(p)

#     score = entail_prob - contra_prob
#     score = (score + 1) / 2  # normalize [-1,1] → [0,1]
#     return round(max(0.0, min(1.0, score)), 3)


# def check_assumption(response: str) -> float:
#     """Detect speculative/hedging terms."""
#     if not response:
#         return 0.0
#     speculative_terms = ["maybe", "probably", "might", "perhaps", "i guess", "seems", "could"]
#     count = sum(1 for t in speculative_terms if t in response.lower())
#     score = 1.0 - min(count / 5.0, 1.0)  # smoother decay
#     return round(score, 3)


# def check_coherence(response: str) -> float:
#     """Heuristic coherence metric: penalizes very short/long, rewards sentence balance."""
#     if not response:
#         return 0.0
#     words = len(re.findall(r"\w+", response))
#     sents = max(1, len(re.split(r"[.!?]+", response)) - 1)
#     if words < 5:
#         return 0.3
#     if words > 200:
#         return 0.5
#     base = min(1.0, (words / 50.0) + (sents / 5.0))
#     return round(max(0.4, min(base, 0.95)), 3)


# def check_accuracy(reference: str, response: str) -> float:
#     """Semantic similarity between reference and response via embeddings (cosine)."""
#     if not reference or not response:
#         return 0.0
#     ref_emb = embed_model.encode(reference, convert_to_tensor=True)
#     resp_emb = embed_model.encode(response, convert_to_tensor=True)
#     sim = float(util.cos_sim(ref_emb, resp_emb).item())
#     return round(max(0.0, min(1.0, sim)), 3)


# # --------------------------
# # ROW & DF EVALUATION
# # --------------------------
# def evaluate_row(row):
#     prompt = row.get("prompt", "")
#     response = row.get("response", "")
#     reference = row.get("reference", "")

#     metrics = {
#         "task_id": row.get("task_id", ""),
#         "agent": row.get("agent", ""),
#         "instruction_following": check_instruction_following(prompt, response),
#         "hallucination": check_hallucination(reference, response),
#         "assumption": check_assumption(response),
#         "coherence": check_coherence(response),
#         "accuracy": check_accuracy(reference, response),
#     }

#     # Weighted avg score (you can adjust weights)
#     metrics["final_score"] = round(
#         0.25 * metrics["instruction_following"]
#         + 0.25 * metrics["accuracy"]
#         + 0.2 * metrics["hallucination"]
#         + 0.15 * metrics["coherence"]
#         + 0.15 * metrics["assumption"],
#         3,
#     )
#     return metrics


# def evaluate_dataframe(df: pd.DataFrame):
#     metrics_df = df.apply(evaluate_row, axis=1, result_type="expand")

#     # Leaderboard
#     leaderboard = (
#         metrics_df.groupby(["agent", "task_id"])["final_score"]
#         .mean()
#         .reset_index()
#     )


#     # # Plots
#     # images = []
#     # Existing images list
#     images = []
    
#     # Add radar chart
#     radar_path = plot_radar_chart(metrics_df, agents=df["agent"].unique(), 
#
###############################################################################################################################

# # evaluator.py
# import math
# import uuid
# import json
# import re
# from typing import List, Dict, Tuple
# import numpy as np
# import pandas as pd
# import matplotlib.pyplot as plt
# import seaborn as sns

# # -----------------------------
# # Lazy model loading
# # -----------------------------
# NLI_MODEL = "textattack/roberta-base-MNLI"
# EMBED_MODEL = "sentence-transformers/all-MiniLM-L6-v2"

# _nli_tokenizer = None
# _nli_model = None
# _embed_model = None
# _id2label = None

# def ensure_models_loaded():
#     global _nli_tokenizer, _nli_model, _embed_model, _id2label, util
#     if _embed_model is None:
#         from sentence_transformers import SentenceTransformer, util
#         _embed_model = SentenceTransformer(EMBED_MODEL)
#         globals()["util"] = util
#     if _nli_model is None:
#         from transformers import AutoTokenizer, AutoModelForSequenceClassification
#         _nli_tokenizer = AutoTokenizer.from_pretrained(NLI_MODEL)
#         _nli_model = AutoModelForSequenceClassification.from_pretrained(NLI_MODEL)
#         _nli_model.to("cpu")
#         _nli_model.eval()
#         _id2label = {int(k): v.upper() for k, v in _nli_model.config.id2label.items()}

# def get_embed_model():
#     ensure_models_loaded()
#     return _embed_model

# def get_nli_tokenizer_and_model():
#     ensure_models_loaded()
#     return _nli_tokenizer, _nli_model, _id2label


# # -----------------------------
# # Metric functions
# # -----------------------------

# def check_instruction_following(prompt: str, response: str) -> float:
#     if not prompt or not response:
#         return 0.0
#     embed_model = get_embed_model()
#     p_emb = embed_model.encode(prompt, convert_to_tensor=True)
#     r_emb = embed_model.encode(response, convert_to_tensor=True)
#     sim = float(util.cos_sim(p_emb, r_emb).item())
#     return round(max(0.0, min(1.0, sim)), 3)

# def check_hallucination(prompt: str, response: str) -> float:
#     if not prompt or not response:
#         return 0.0
#     tokenizer, model, id2label = get_nli_tokenizer_and_model()
#     inputs = tokenizer.encode_plus(prompt, response, return_tensors="pt", truncation=True)
#     outputs = model(**inputs)
#     probs = outputs.logits.softmax(dim=1).detach().cpu().numpy()[0]
#     labels = [id2label[i] for i in range(len(probs))]
#     if "ENTAILMENT" in labels:
#         entailment_prob = float(probs[labels.index("ENTAILMENT")])
#     else:
#         entailment_prob = float(probs.max())
#     return round(entailment_prob, 3)

# def check_accuracy(prompt: str, response: str, metadata: str) -> float:
#     """
#     Accuracy:
#       - if metadata contains an 'answer', compare response vs answer using embeddings
#       - else fall back to entailment probability (NLI)
#     """
#     answer = None
#     if isinstance(metadata, str) and metadata.strip().startswith("{"):
#         try:
#             meta = json.loads(metadata)
#             for k in ("answer", "reference", "ground_truth"):
#                 if k in meta:
#                     answer = meta[k]
#                     break
#         except Exception:
#             pass
#     elif isinstance(metadata, dict):
#         for k in ("answer", "reference", "ground_truth"):
#             if k in metadata:
#                 answer = metadata[k]
#                 break

#     if answer:
#         embed_model = get_embed_model()
#         ref_emb = embed_model.encode(str(answer), convert_to_tensor=True)
#         r_emb = embed_model.encode(response, convert_to_tensor=True)
#         sim = float(util.cos_sim(ref_emb, r_emb).item())
#         return round(max(0.0, min(1.0, sim)), 3)
#     else:
#         return check_hallucination(prompt, response)

# def check_coherence(response: str) -> float:
#     """
#     Coherence: average pairwise sentence similarity if >1 sentence,
#     otherwise fallback heuristic.
#     """
#     if not response or not isinstance(response, str) or response.strip() == "":
#         return 0.25
#     sents = [s.strip() for s in re.split(r"[.!?]+", response) if s.strip()]
#     embed_model = get_embed_model()
#     if len(sents) > 1:
#         try:
#             embs = embed_model.encode(sents, convert_to_tensor=True)
#             sims = []
#             for i in range(len(embs)):
#                 for j in range(i + 1, len(embs)):
#                     sims.append(float(util.cos_sim(embs[i], embs[j]).item()))
#             if sims:
#                 score = (np.mean(sims) + 1.0) / 2.0
#                 return round(max(0.0, min(1.0, score)), 3)
#         except Exception:
#             pass
#     words = len(re.findall(r"\w+", response))
#     return round(min(0.8, max(0.3, words / 50.0)), 3)

# def check_fluency(response: str) -> float:
#     if not response or not isinstance(response, str):
#         return 0.0
#     letters = sum(ch.isalpha() for ch in response)
#     total = len(response)
#     return round(max(0.0, min(1.0, letters / max(1, total))), 3)


# # -----------------------------
# # Visualization helpers
# # -----------------------------

# def spider_net_multi(labels: List[str], rows: List[Dict], title: str, fill_alpha: float = 0.12):
#     N = len(labels)
#     angles = [n / float(N) * 2 * math.pi for n in range(N)]
#     angles += angles[:1]
#     fig = plt.figure(figsize=(6.5, 6.5))
#     ax = plt.subplot(111, polar=True)
#     ax.set_xticks(angles[:-1])
#     ax.set_xticklabels(labels, fontsize=9)
#     ax.set_ylim(0, 1)
#     ax.set_yticks([0, 0.25, 0.5, 0.75, 1.0])
#     for r in rows:
#         values = r["values"]
#         values_closed = values + values[:1]
#         ax.plot(angles, values_closed, linewidth=1.5, label=r["name"])
#         ax.fill(angles, values_closed, alpha=fill_alpha)
#     ax.set_title(title, y=1.08, fontsize=12)
#     ax.legend(loc="upper right", bbox_to_anchor=(1.25, 1.1))
#     return fig

# def heatmap_plot(df: pd.DataFrame, metric_cols: List[str], title: str = "Metric Correlations"):
#     fig, ax = plt.subplots(figsize=(7, 5))
#     sns.heatmap(df[metric_cols].corr(), annot=True, fmt=".2f", cmap="coolwarm", ax=ax)
#     ax.set_title(title)
#     return fig


# # -----------------------------
# # Main evaluation entrypoint
# # -----------------------------

# def evaluate_dataframe(df: pd.DataFrame) -> Tuple[pd.DataFrame, List[Tuple[str,str]], pd.DataFrame]:
#     df = df.copy()

#     scores = []
#     for _, r in df.iterrows():
#         prompt = str(r.get("prompt", "") or "")
#         response = str(r.get("response", "") or "")
#         metadata = r.get("metadata", "")

#         s = {}
#         s["instruction_following"] = check_instruction_following(prompt, response)
#         s["hallucination"] = check_hallucination(prompt, response)
#         s["accuracy"] = check_accuracy(prompt, response, metadata)
#         s["coherence"] = check_coherence(response)
#         s["fluency"] = check_fluency(response)

#         # clamp and nanmean
#         for k in s:
#             try:
#                 s[k] = float(s[k])
#             except Exception:
#                 s[k] = np.nan
#             if not np.isnan(s[k]):
#                 s[k] = max(0.0, min(1.0, s[k]))

#         vals = [v for v in s.values() if not np.isnan(v)]
#         s["FinalScore"] = round(float(np.nanmean(vals)) if vals else 0.0, 3)
#         scores.append(s)

#     metrics_df = pd.concat([df.reset_index(drop=True), pd.DataFrame(scores)], axis=1)

#     metric_cols = ["instruction_following", "hallucination", "accuracy", "coherence", "fluency", "FinalScore"]
#     leaderboard = (
#         metrics_df.groupby(["agent", "task_type"])[metric_cols]
#         .mean()
#         .reset_index()
#     )

#     images = []

#     try:
#         rows = []
#         for agent in leaderboard["agent"].unique():
#             vals = leaderboard[leaderboard["agent"] == agent][metric_cols[:-1]].mean().tolist()
#             rows.append({"name": agent, "values": vals})
#         fig1 = spider_net_multi(metric_cols[:-1], rows, "Agent Performance Radar")
#         path1 = f"/tmp/radar_{uuid.uuid4().hex}.png"
#         fig1.savefig(path1, bbox_inches="tight")
#         plt.close(fig1)
#         images.append((path1, "Radar Plot"))
#     except Exception as e:
#         print("Radar plot failed:", e)

#     try:
#         fig2 = heatmap_plot(metrics_df, metric_cols[:-1], title="Metric Correlation Heatmap")
#         path2 = f"/tmp/heatmap_{uuid.uuid4().hex}.png"
#         fig2.savefig(path2, bbox_inches="tight")
#         plt.close(fig2)
#         images.append((path2, "Correlation Heatmap"))
#     except Exception as e:
#         print("Heatmap failed:", e)

#     return metrics_df, images, leaderboard

# evaluator.py
"""
Evaluator module for agentic-evaluation-framework.

Provides:
 - Evaluator class that computes:
   - BERTScore (if reference provided)
   - Semantic similarity (SBERT)
   - NLI entailment (roberta-large-mnli)
   - Optional LLM-as-a-judge scoring (flan-t5 or other seq2seq)
 - Aggregation per agent, ranking, and visualizations (plotly)
"""

import json
import re
import math
from typing import Optional, Dict, Any, List
import numpy as np
import pandas as pd
from tqdm.auto import tqdm

import torch
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    AutoModelForSeq2SeqLM,
    pipeline,
)
from sentence_transformers import SentenceTransformer
from bert_score import score as bert_score
from sklearn.metrics.pairwise import cosine_similarity

import plotly.express as px
import plotly.graph_objects as go

# Helper
def _safe_json_load(text: str):
    """Try to extract JSON object from a model's output robustly."""
    text = text.strip()
    # Try direct JSON
    try:
        return json.loads(text)
    except Exception:
        pass
    # Try to find {...}
    m = re.search(r"(\{.*\})", text, flags=re.S)
    if m:
        try:
            return json.loads(m.group(1))
        except Exception:
            pass
    # Try python-literal fallback
    try:
        import ast
        return ast.literal_eval(text)
    except Exception:
        return None

class Evaluator:
    def __init__(
        self,
        judge_model_name: Optional[str] = None,
        device: Optional[str] = None,
        embedding_model_name: str = "sentence-transformers/all-MiniLM-L6-v2",
        nli_model_name: str = "roberta-large-mnli",
        bertscore_model: str = "roberta-large"
    ):
        """
        judge_model_name: (optional) seq2seq model (e.g. "google/flan-t5-base") used as LLM judge
        device: "cpu" or "cuda" or None (auto)
        """
        # Device detection
        if device is None:
            self.device = "cuda" if torch.cuda.is_available() else "cpu"
        else:
            self.device = device

        # Sentence embedding model (light and fast)
        self.embed = SentenceTransformer(embedding_model_name, device=self.device)

        # NLI model (for entailment-based hallucination signals)
        self.nli_tokenizer = AutoTokenizer.from_pretrained(nli_model_name)
        self.nli_model = AutoModelForSequenceClassification.from_pretrained(nli_model_name)
        self.nli_model.to(self.device)
        # label mapping from config
        self.nli_id2label = {int(k): v for k, v in self.nli_model.config.id2label.items()}

        self.bertscore_model = bertscore_model

        # Optional judge model (text2text seq2seq)
        self.judge_model_name = judge_model_name
        self.judge_pipe = None
        if judge_model_name and judge_model_name.lower() != "none":
            # Use text2text-generation pipeline
            gen_device = 0 if (self.device == "cuda") else -1
            self.judge_tokenizer = AutoTokenizer.from_pretrained(judge_model_name)
            self.judge_model = AutoModelForSeq2SeqLM.from_pretrained(judge_model_name)
            # Move judge to device if possible
            if self.device == "cuda":
                self.judge_model.to(self.device)
                gen_device = 0
            self.judge_pipe = pipeline(
                "text2text-generation",
                model=self.judge_model,
                tokenizer=self.judge_tokenizer,
                device=gen_device,
                # we set deterministic decode later per call
            )

    # -------------------------
    # Core metric funcs
    # -------------------------
    def _semantic_similarity(self, ref: str, resp: str) -> float:
        """SBERT cosine similarity in [0,1]."""
        emb = self.embed.encode([ref, resp], convert_to_numpy=True)
        sim = float(cosine_similarity([emb[0]], [emb[1]])[0][0])
        # Clamp
        return max(0.0, min(1.0, sim))

    def _bertscore_f1(self, refs: List[str], hyps: List[str]) -> float:
        """Compute BERTScore F1 (returns average F1). Uses bert_score package."""
        try:
            P, R, F1 = bert_score(hyps, refs, model_type=self.bertscore_model, lang="en", rescale_with_baseline=True)
            # F1 is a torch tensor array
            return float(F1.mean().item())
        except Exception as e:
            # fallback: 0.0 if too heavy or failed
            print("BERTScore error:", e)
            return 0.0

    def _nli_entailment_prob(self, premise: str, hypothesis: str) -> float:
        """
        Returns probability of ENTAILMENT in [0,1] for (premise=reference, hypothesis=response).
        Uses the loaded NLI model.
        """
        self.nli_model.eval()
        tok = self.nli_tokenizer.encode_plus(
            premise,
            hypothesis,
            return_tensors="pt",
            truncation=True,
            max_length=512,
        )
        inputs = {k: v.to(self.device) for k, v in tok.items()}
        with torch.no_grad():
            out = self.nli_model(**inputs)
            logits = out.logits[0]  # shape (3,)
            probs = torch.softmax(logits, dim=-1).detach().cpu().numpy()
        # Find label index for entailment (case-insensitive)
        entail_idx = None
        for k, v in self.nli_id2label.items():
            if "ENTAIL" in v.upper():
                entail_idx = k
                break
        if entail_idx is None:
            # fallback assume label 2 is entailment
            entail_idx = 2 if len(probs) > 2 else np.argmax(probs)
        return float(probs[entail_idx])

    # -------------------------
    # LLM-as-judge (rubric)
    # -------------------------
    def _llm_judge_single(self, prompt: str, response: str, reference: Optional[str] = None) -> Dict[str, Any]:
        """
        Ask judge model to evaluate single example and return parsed JSON:
        {
          "instruction_following": int(0-100),
          "hallucination": int(0-100),        # 0 = no hallucination, 100 = extreme hallucination
          "assumption_control": int(0-100),   # 0 = poor, 100 = excellent
          "coherence": int(0-100),
          "accuracy": int(0-100),
          "explanations": { ... }
        }
        """
        if self.judge_pipe is None:
            raise RuntimeError("Judge model not loaded")

        # Build the instruction carefully
        payload = {
            "prompt": prompt,
            "response": response,
            "reference": reference if reference is not None else "",
        }
        instruction = (
            "You are a strict, unbiased evaluator. Given the user PROMPT, the assistant RESPONSE, "
            "and an optional REFERENCE (ground-truth), evaluate the response on five axes:\n\n"
            "1) instruction_following (0-100): how well the response follows the user's instruction.\n"
            "2) hallucination (0-100): how much the response contains ungrounded or fabricated facts. "
            "0 means no hallucination; 100 means severe hallucination.\n"
            "3) assumption_control (0-100): how careful the assistant was about making unsupported assumptions, "
            "or whether it asked clarifying questions when needed. 100 = excellent assumption control.\n"
            "4) coherence (0-100): internal logical flow and clarity.\n"
            "5) accuracy (0-100): factual correctness relative to the reference if present; otherwise, use internal consistency and plausibility.\n\n"
            "Return only a JSON object with integer scores 0..100 and a short explanation text for each axis, "
            "for example:\n"
            '{"instruction_following": 85, "hallucination": 10, "assumption_control": 70, "coherence": 80, "accuracy": 75, "explanations": {"instruction_following": "...", ...}}\n\n'
            f"Input JSON: {json.dumps(payload)}\nOutput (JSON only):"
        )

        # Call the model deterministically
        gen = self.judge_pipe(
            instruction,
            max_length=512,
            do_sample=False,
            num_beams=1,
            return_full_text=False,
        )
        raw = gen[0]["generated_text"] if isinstance(gen, list) else gen["generated_text"]
        parsed = _safe_json_load(raw)
        if parsed is None:
            # If parsing fails, return a conservative default with explanation
            return {
                "instruction_following": 50,
                "hallucination": 50,
                "assumption_control": 50,
                "coherence": 50,
                "accuracy": 50,
                "explanations": {"error": "Judge output parsing failed", "raw": raw},
            }
        # Normalize numeric values and ensure keys exist
        for k in ["instruction_following", "hallucination", "assumption_control", "coherence", "accuracy"]:
            if k not in parsed:
                parsed[k] = 50
            else:
                try:
                    parsed[k] = int(round(float(parsed[k])))
                    parsed[k] = max(0, min(100, parsed[k]))
                except Exception:
                    parsed[k] = 50
        if "explanations" not in parsed:
            parsed["explanations"] = {}
        return parsed

    # -------------------------
    # Main evaluation methods
    # -------------------------
    def evaluate_dataframe(self, df: pd.DataFrame, mapping: Dict[str, str], progress: bool = True) -> pd.DataFrame:
        """
        mapping: dict with keys: 'prompt','response','agent','reference' (reference optional), 'task_type','metadata'
        Returns a DataFrame with per-row metrics appended.
        """
        rows = []
        it = df.iterrows()
        if progress:
            it = tqdm(df.iterrows(), total=len(df), desc="Evaluating rows")
        for idx, row in it:
            prompt = row.get(mapping.get("prompt"))
            response = row.get(mapping.get("response"))
            agent = row.get(mapping.get("agent"), "unknown-agent")
            reference = None
            if mapping.get("reference") and mapping.get("reference") in row:
                reference = row.get(mapping.get("reference"))
            # Compute metrics
            r = {
                "index": int(idx),
                "agent": agent,
                "prompt": prompt,
                "response": response,
                "reference": reference,
            }
            # Reference-based if present
            if reference and isinstance(reference, str) and len(reference.strip()) > 0:
                try:
                    r["sbert_sim"] = self._semantic_similarity(reference, response)
                except Exception as e:
                    r["sbert_sim"] = None
                try:
                    r["bertscore_f1"] = self._bertscore_f1([reference], [response])
                except Exception:
                    r["bertscore_f1"] = None
                try:
                    r["nli_entail_prob"] = self._nli_entailment_prob(reference, response)
                except Exception:
                    r["nli_entail_prob"] = None
            else:
                # No reference: still compute sbert similarity to prompt (weak proxy)
                try:
                    r["sbert_sim"] = self._semantic_similarity(prompt if isinstance(prompt, str) else "", response)
                except Exception:
                    r["sbert_sim"] = None
                r["bertscore_f1"] = None
                r["nli_entail_prob"] = None

            # LLM judge (if available)
            if self.judge_pipe is not None:
                try:
                    judge_out = self._llm_judge_single(prompt or "", response or "", reference)
                    # store judge outputs
                    r.update({
                        "jf_instruction_following": judge_out.get("instruction_following"),
                        "jf_hallucination": judge_out.get("hallucination"),
                        "jf_assumption_control": judge_out.get("assumption_control"),
                        "jf_coherence": judge_out.get("coherence"),
                        "jf_accuracy": judge_out.get("accuracy"),
                        "jf_explanations": json.dumps(judge_out.get("explanations", {}))[:2000],
                    })
                except Exception as e:
                    print("Judge error:", e)
                    r.update({
                        "jf_instruction_following": None,
                        "jf_hallucination": None,
                        "jf_assumption_control": None,
                        "jf_coherence": None,
                        "jf_accuracy": None,
                        "jf_explanations": None,
                    })
            else:
                # No judge -> leave judge fields None
                r.update({
                    "jf_instruction_following": None,
                    "jf_hallucination": None,
                    "jf_assumption_control": None,
                    "jf_coherence": None,
                    "jf_accuracy": None,
                    "jf_explanations": None,
                })
            rows.append(r)
        results_df = pd.DataFrame(rows)
        return results_df

    # -------------------------
    # Aggregation + scoring
    # -------------------------
    def aggregate_agent_scores(
        self,
        per_example_df: pd.DataFrame,
        use_judge_if_available: bool = True,
        weights: Optional[Dict[str, float]] = None
    ) -> pd.DataFrame:
        """
        Aggregates per-example metrics to per-agent scores and computes overall_score.
        If judge outputs are available and use_judge_if_available=True, prefer those.
        weights: mapping of metric->weight (must sum to 1). Supported keys:
            instruction_following, hallucination (lower is better), assumption_control, coherence, accuracy
        """
        if weights is None:
            weights = {
                "instruction_following": 0.25,
                "hallucination": 0.20,
                "assumption_control": 0.10,
                "coherence": 0.20,
                "accuracy": 0.25
            }
        agents = []
        for agent, g in per_example_df.groupby("agent"):
            agent_row = {"agent": agent, "n_examples": len(g)}
            # Decide metrics: prefer judge fields
            if use_judge_if_available and g["jf_instruction_following"].notna().any():
                # use judge means take mean of jf_* values across samples (0..100)
                inst = float(g["jf_instruction_following"].dropna().mean())
                hall = float(g["jf_hallucination"].dropna().mean())
                assum = float(g["jf_assumption_control"].dropna().mean())
                coh = float(g["jf_coherence"].dropna().mean())
                acc = float(g["jf_accuracy"].dropna().mean())
            else:
                # Fallback heuristics (less reliable):
                # instruction_following ~ berts_f1 (0..1) scaled to 0..100
                inst = float(g["bertscore_f1"].dropna().mean() * 100) if g["bertscore_f1"].notna().any() else float(g["sbert_sim"].dropna().mean() * 100 if g["sbert_sim"].notna().any() else 50.0)
                # hallucination ~ (1 - nli_entail_prob) * 100 if available, else 50
                if g["nli_entail_prob"].notna().any():
                    hall = float((1.0 - g["nli_entail_prob"].dropna().mean()) * 100.0)
                else:
                    hall = 50.0
                # assumption_control: unknown without judge -> default 50
                assum = 50.0
                # coherence: approximate with sbert_sim to prompt (0..1)
                coh = float(g["sbert_sim"].dropna().mean() * 100) if g["sbert_sim"].notna().any() else 50.0
                # accuracy: prefer sbert_sim or berts_f1
                if g["sbert_sim"].notna().any():
                    acc = float(g["sbert_sim"].dropna().mean() * 100.0)
                elif g["bertscore_f1"].notna().any():
                    acc = float(g["bertscore_f1"].dropna().mean() * 100.0)
                else:
                    acc = 50.0
            # store
            agent_row.update({
                "instruction_following": inst,
                "hallucination": hall,
                "assumption_control": assum,
                "coherence": coh,
                "accuracy": acc,
            })
            # compute overall score (hallucination is negative -> convert to positive reliability = 100 - hallucination)
            overall = (
                weights["instruction_following"] * inst +
                weights["hallucination"] * (100.0 - hall) +
                weights["assumption_control"] * assum +
                weights["coherence"] * coh +
                weights["accuracy"] * acc
            )
            agent_row["overall_score"] = overall
            agents.append(agent_row)
        agg = pd.DataFrame(agents)
        agg = agg.sort_values("overall_score", ascending=False).reset_index(drop=True)
        return agg

    # -------------------------
    # Visualizations (plotly)
    # -------------------------
    def plot_leaderboard(self, agg_df: pd.DataFrame, top_k: int = 10):
        df = agg_df.head(top_k).copy()
        fig = px.bar(df, x="agent", y="overall_score", color="overall_score",
                     labels={"overall_score":"Overall score"}, title=f"Leaderboard (top {top_k})")
        fig.update_layout(yaxis_range=[0,100])
        return fig

    def plot_metric_heatmap(self, agg_df: pd.DataFrame):
        metrics = ["instruction_following", "hallucination", "assumption_control", "coherence", "accuracy"]
        hm = agg_df.set_index("agent")[metrics].copy()
        # For hallucination, we prefer to show "hallucination (lower better)" as-is; we'll annotate colorbar accordingly
        fig = px.imshow(hm, labels=dict(x="metric", y="agent", color="value"), aspect="auto", title="Agent × Metric Heatmap")
        return fig

    def plot_radar(self, agg_df: pd.DataFrame, agents_to_plot: Optional[List[str]] = None):
        metrics = ["instruction_following", "hallucination", "assumption_control", "coherence", "accuracy"]
        # convert hallucination to 100-hallucination so larger is better for radar
        df = agg_df.copy()
        df["hallucination_reliability"] = 100.0 - df["hallucination"]
        plot_metrics = ["instruction_following", "hallucination_reliability", "assumption_control", "coherence", "accuracy"]
        if agents_to_plot is None:
            agents_to_plot = df["agent"].head(5).tolist()
        fig = go.Figure()
        for a in agents_to_plot:
            row = df[df["agent"] == a].iloc[0]
            values = [row[m] for m in plot_metrics]
            fig.add_trace(go.Scatterpolar(r=values, theta=plot_metrics, fill='toself', name=a))
        fig.update_layout(polar=dict(radialaxis=dict(visible=True, range=[0,100])), showlegend=True, title="Radar: per-agent metrics (higher better)")
        return fig