Agent Evaluation and Benchmarking

def evaluate_with_variance(
    agent_fn, test_case: dict, num_runs: int = 5
) -> dict:
    """Run an agent multiple times and report variance.

    Args:
        agent_fn: The agent function to evaluate.
        test_case: Dict with 'input' and 'expected_output' keys.
        num_runs: Number of independent runs.

    Returns:
        Dict with mean score, standard deviation, and individual results.
    """
    scores = []
    outputs = []

    for run in range(num_runs):
        output = agent_fn(test_case["input"])
        score = compute_score(output, test_case["expected_output"])
        scores.append(score)
        outputs.append(output)

    return {
        "mean_score": sum(scores) / len(scores),
        "std_dev": (sum((s - sum(scores)/len(scores))**2 for s in scores) / len(scores)) ** 0.5,
        "min_score": min(scores),
        "max_score": max(scores),
        "num_runs": num_runs,
        "all_scores": scores,
    }

def task_success_rate(
    results: list[dict],
) -> dict:
    """Compute task-level success metrics.

    Args:
        results: List of dicts with 'task_id', 'completed', and 'correct' fields.

    Returns:
        Dict with success metrics.
    """
    n = len(results)
    completed = sum(1 for r in results if r["completed"])
    correct = sum(1 for r in results if r["correct"])

    return {
        "total_tasks": n,
        "completion_rate": completed / n if n > 0 else 0.0,
        "success_rate": correct / n if n > 0 else 0.0,
        "completed_but_incorrect": completed - correct,
    }

from dataclasses import dataclass


@dataclass
class TaskResult:
    """Result of a single task evaluation."""

    task_id: str
    score: float          # 0.0 to 1.0 (partial credit)
    completed: bool       # Did the agent declare completion?
    steps_taken: int      # Number of steps to reach the result
    errors_encountered: int
    tools_used: list[str]
    wall_time_seconds: float
    tokens_consumed: int
    notes: str = ""


def evaluate_task_quality(
    result: TaskResult,
    max_steps: int = 20,
    max_time: float = 300.0,
) -> dict:
    """Evaluate task completion quality with multiple dimensions.

    Combines correctness with efficiency metrics.
    """
    # Correctness (primary metric)
    correctness = result.score

    # Efficiency: penalize excessive steps
    step_efficiency = max(0.0, 1.0 - (result.steps_taken / max_steps))

    # Time efficiency
    time_efficiency = max(0.0, 1.0 - (result.wall_time_seconds / max_time))

    # Error rate (lower is better)
    error_rate = result.errors_encountered / max(result.steps_taken, 1)

    # Combined quality score (weighted)
    quality_score = (
        0.6 * correctness +
        0.2 * step_efficiency +
        0.1 * time_efficiency +
        0.1 * (1 - error_rate)
    )

    return {
        "task_id": result.task_id,
        "correctness": correctness,
        "step_efficiency": step_efficiency,
        "time_efficiency": time_efficiency,
        "error_rate": error_rate,
        "quality_score": quality_score,
    }

CODING_TASK_RUBRIC = {
    "correctness": {
        "description": "Does the code produce the correct output?",
        "weight": 0.40,
        "levels": {
            1.0: "All test cases pass",
            0.75: "Most test cases pass (>80%)",
            0.5: "Some test cases pass (40-80%)",
            0.25: "Few test cases pass (<40%)",
            0.0: "No test cases pass or code does not run",
        },
    },
    "code_quality": {
        "description": "Is the code clean, readable, and well-structured?",
        "weight": 0.20,
        "levels": {
            1.0: "Excellent: clean, well-documented, good variable names",
            0.75: "Good: mostly clean with minor issues",
            0.5: "Acceptable: works but has style issues",
            0.25: "Poor: hard to read, bad structure",
            0.0: "Very poor: incomprehensible",
        },
    },
    "efficiency": {
        "description": "Is the solution reasonably efficient?",
        "weight": 0.15,
        "levels": {
            1.0: "Optimal time and space complexity",
            0.75: "Near-optimal",
            0.5: "Acceptable for the problem size",
            0.25: "Unnecessarily slow or memory-intensive",
            0.0: "Unacceptably inefficient",
        },
    },
    "error_handling": {
        "description": "Does the code handle edge cases and errors?",
        "weight": 0.15,
        "levels": {
            1.0: "Comprehensive error handling",
            0.5: "Basic error handling",
            0.0: "No error handling",
        },
    },
    "process_quality": {
        "description": "Was the agent's problem-solving process reasonable?",
        "weight": 0.10,
        "levels": {
            1.0: "Systematic: understood, planned, implemented, tested",
            0.5: "Somewhat systematic but with unnecessary steps",
            0.0: "Chaotic: random trial and error",
        },
    },
}


def rubric_evaluate(scores: dict[str, float], rubric: dict) -> float:
    """Compute a weighted score from rubric evaluations.

    Args:
        scores: Dict mapping rubric criterion names to scores (0.0-1.0).
        rubric: The rubric definition with criteria and weights.

    Returns:
        Weighted total score.
    """
    total = 0.0
    total_weight = 0.0
    for criterion, config in rubric.items():
        if criterion in scores:
            total += scores[criterion] * config["weight"]
            total_weight += config["weight"]
    return total / total_weight if total_weight > 0 else 0.0

REASONING_EVALUATION_PROMPT = """Evaluate the quality of the following agent's
reasoning process. Score each dimension from 1-5.

Agent's reasoning trace:
{trace}

Correct answer: {correct_answer}

Dimensions to evaluate:
1. **Logical Coherence** (1-5): Does each step follow logically from the previous?
2. **Relevance** (1-5): Are the reasoning steps relevant to the problem?
3. **Completeness** (1-5): Does the reasoning address all aspects of the problem?
4. **Efficiency** (1-5): Is the reasoning concise without unnecessary tangents?
5. **Self-Correction** (1-5): When errors occur, does the agent recognize and fix them?

For each dimension, provide a score and brief justification.
"""


def evaluate_reasoning(
    trace: str, correct_answer: str, llm_call
) -> dict:
    """Evaluate reasoning quality using an LLM judge.

    Args:
        trace: The agent's full reasoning trace.
        correct_answer: The known correct answer.
        llm_call: Function to call the judge LLM.

    Returns:
        Dict with scores for each reasoning dimension.
    """
    response = llm_call(
        prompt=REASONING_EVALUATION_PROMPT.format(
            trace=trace, correct_answer=correct_answer
        )
    )
    return parse_dimension_scores(response)


def parse_dimension_scores(response: str) -> dict:
    """Parse dimension scores from judge response."""
    dimensions = [
        "logical_coherence",
        "relevance",
        "completeness",
        "efficiency",
        "self_correction",
    ]
    scores = {}
    for dim in dimensions:
        # Look for patterns like "Logical Coherence: 4/5" or "Score: 4"
        for line in response.split("\n"):
            dim_readable = dim.replace("_", " ")
            if dim_readable.lower() in line.lower() and any(c.isdigit() for c in line):
                digits = [int(c) for c in line if c.isdigit()]
                if digits:
                    scores[dim] = digits[0] / 5.0  # Normalize to 0-1
                    break
        if dim not in scores:
            scores[dim] = 0.5  # Default if parsing fails
    return scores

def evaluate_tool_use(
    tool_calls: list[dict], optimal_calls: list[dict] | None = None
) -> dict:
    """Evaluate the efficiency and appropriateness of tool usage.

    Args:
        tool_calls: List of actual tool calls the agent made.
        optimal_calls: List of ideal tool calls (if known).

    Returns:
        Dict with tool use metrics.
    """
    total_calls = len(tool_calls)

    # Count redundant calls (same tool with same/similar arguments)
    seen = set()
    redundant = 0
    for call in tool_calls:
        key = f"{call['tool']}:{call.get('args', '')}"
        if key in seen:
            redundant += 1
        seen.add(key)

    # Count failed calls
    failed = sum(1 for call in tool_calls if call.get("status") == "error")

    # Tool diversity (number of unique tools used)
    unique_tools = len({call["tool"] for call in tool_calls})

    metrics = {
        "total_calls": total_calls,
        "unique_tools_used": unique_tools,
        "redundant_calls": redundant,
        "failed_calls": failed,
        "redundancy_rate": redundant / total_calls if total_calls > 0 else 0.0,
        "failure_rate": failed / total_calls if total_calls > 0 else 0.0,
    }

    # Compare to optimal if available
    if optimal_calls is not None:
        optimal_count = len(optimal_calls)
        metrics["optimal_calls"] = optimal_count
        metrics["call_overhead"] = (total_calls - optimal_count) / optimal_count if optimal_count > 0 else 0.0

    return metrics

def analyze_trajectory(
    trajectory: list[dict],
) -> dict:
    """Analyze an agent's complete action trajectory.

    Args:
        trajectory: List of actions, each with 'type', 'content', and 'timestamp'.

    Returns:
        Dict with trajectory analysis metrics.
    """
    total_steps = len(trajectory)

    # Categorize steps
    reasoning_steps = [s for s in trajectory if s["type"] == "thought"]
    action_steps = [s for s in trajectory if s["type"] == "action"]
    observation_steps = [s for s in trajectory if s["type"] == "observation"]

    # Detect backtracking (similar actions repeated)
    backtrack_count = 0
    for i in range(1, len(action_steps)):
        if action_steps[i]["content"] == action_steps[i-1]["content"]:
            backtrack_count += 1

    # Compute timing if timestamps available
    if trajectory and "timestamp" in trajectory[0]:
        from datetime import datetime
        start = datetime.fromisoformat(trajectory[0]["timestamp"])
        end = datetime.fromisoformat(trajectory[-1]["timestamp"])
        duration = (end - start).total_seconds()
    else:
        duration = None

    return {
        "total_steps": total_steps,
        "reasoning_steps": len(reasoning_steps),
        "action_steps": len(action_steps),
        "observation_steps": len(observation_steps),
        "reasoning_to_action_ratio": (
            len(reasoning_steps) / len(action_steps) if action_steps else 0
        ),
        "backtrack_count": backtrack_count,
        "duration_seconds": duration,
    }

JUDGE_PROMPT = """You are an expert judge evaluating the quality of an AI agent's
response to a task.

Task: {task}

Agent's response:
{response}

Evaluate on a scale of 1-10 across these dimensions:
1. **Correctness**: Is the response factually accurate and logically sound?
2. **Completeness**: Does the response fully address all aspects of the task?
3. **Clarity**: Is the response well-organized and easy to understand?
4. **Usefulness**: Would this response be practically helpful to the user?

For each dimension, provide:
- Score (1-10)
- Brief justification

Finally, provide an overall score (1-10).
"""


def llm_as_judge(
    task: str,
    response: str,
    llm_call,
    judge_model: str = "gpt-4",
) -> dict:
    """Use an LLM to evaluate an agent's response.

    Args:
        task: The original task description.
        response: The agent's response to evaluate.
        llm_call: Function to call the judge LLM.
        judge_model: Model to use as judge.

    Returns:
        Dict with dimension scores and overall score.
    """
    prompt = JUDGE_PROMPT.format(task=task, response=response)
    judgment = llm_call(model=judge_model, prompt=prompt)
    return parse_judgment(judgment)


def parse_judgment(judgment: str) -> dict:
    """Parse a judge's response into structured scores."""
    dimensions = ["correctness", "completeness", "clarity", "usefulness"]
    scores = {}

    for dim in dimensions:
        for line in judgment.split("\n"):
            if dim.lower() in line.lower() and any(c.isdigit() for c in line):
                digits = [int(c) for c in line if c.isdigit()]
                if digits:
                    scores[dim] = min(digits[0], 10) / 10.0  # Normalize to 0-1
                    break

    # Extract overall score
    for line in judgment.split("\n"):
        if "overall" in line.lower() and any(c.isdigit() for c in line):
            digits = [int(c) for c in line if c.isdigit()]
            if digits:
                scores["overall"] = min(digits[0], 10) / 10.0
                break

    if "overall" not in scores:
        scores["overall"] = sum(scores.values()) / len(scores) if scores else 0.5

    scores["raw_judgment"] = judgment
    return scores

PAIRWISE_PROMPT = """Compare the following two responses to the same task.

Task: {task}

Response A:
{response_a}

Response B:
{response_b}

Which response is better? Consider correctness, completeness, clarity, and
usefulness.

Respond with:
- "A" if Response A is clearly better
- "B" if Response B is clearly better
- "TIE" if they are roughly equivalent

Then explain your reasoning in 2-3 sentences."""


def pairwise_compare(
    task: str,
    response_a: str,
    response_b: str,
    llm_call,
    num_comparisons: int = 3,
) -> dict:
    """Compare two responses with position debiasing.

    Runs the comparison multiple times, swapping positions to reduce
    position bias (LLMs tend to prefer the first response).

    Args:
        task: The original task.
        response_a: First response.
        response_b: Second response.
        llm_call: Function to call the judge LLM.
        num_comparisons: Number of comparisons (should be odd).

    Returns:
        Dict with the winner and comparison details.
    """
    a_wins = 0
    b_wins = 0
    ties = 0

    for i in range(num_comparisons):
        # Alternate positions to reduce position bias
        if i % 2 == 0:
            first, second = response_a, response_b
            label_first, label_second = "A", "B"
        else:
            first, second = response_b, response_a
            label_first, label_second = "B", "A"

        prompt = PAIRWISE_PROMPT.format(
            task=task, response_a=first, response_b=second
        )
        result = llm_call(prompt=prompt).strip()

        if result.startswith("A"):
            if label_first == "A":
                a_wins += 1
            else:
                b_wins += 1
        elif result.startswith("B"):
            if label_second == "A":
                a_wins += 1
            else:
                b_wins += 1
        else:
            ties += 1

    if a_wins > b_wins:
        winner = "A"
    elif b_wins > a_wins:
        winner = "B"
    else:
        winner = "TIE"

    return {
        "winner": winner,
        "a_wins": a_wins,
        "b_wins": b_wins,
        "ties": ties,
        "comparisons": num_comparisons,
    }

"""
An automated evaluation pipeline for AI agents.

This pipeline:
1. Loads test cases from a dataset
2. Runs the agent on each test case (with retry and timeout)
3. Evaluates results using multiple metrics
4. Generates a comprehensive report
"""

import json
import time
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path


@dataclass
class TestCase:
    """A single test case for agent evaluation."""

    id: str
    input: str
    expected_output: str | None = None
    metadata: dict = field(default_factory=dict)
    difficulty: str = "medium"
    category: str = "general"


@dataclass
class EvaluationResult:
    """Result of evaluating an agent on a test case."""

    test_case_id: str
    agent_output: str
    scores: dict[str, float]
    trajectory: list[dict]
    wall_time_seconds: float
    tokens_used: int
    error: str | None = None


class EvaluationPipeline:
    """Automated evaluation pipeline for AI agents.

    Supports:
    - Multiple evaluation metrics (exact match, LLM judge, custom)
    - Retry logic with timeout
    - Parallel evaluation (optional)
    - Comprehensive reporting
    """

    def __init__(
        self,
        agent_fn,
        evaluators: dict[str, callable],
        timeout_seconds: float = 300.0,
        max_retries: int = 2,
    ):
        self.agent_fn = agent_fn
        self.evaluators = evaluators
        self.timeout_seconds = timeout_seconds
        self.max_retries = max_retries
        self.results: list[EvaluationResult] = []

    def load_test_cases(self, path: str) -> list[TestCase]:
        """Load test cases from a JSON file.

        Expected format:
        [
            {
                "id": "test_001",
                "input": "...",
                "expected_output": "...",
                "difficulty": "easy",
                "category": "coding"
            },
            ...
        ]
        """
        with open(path) as f:
            data = json.load(f)
        return [TestCase(**case) for case in data]

    def run_agent(self, test_case: TestCase) -> dict:
        """Run the agent on a single test case with timing and error handling."""
        start_time = time.time()
        trajectory = []
        error = None
        output = ""

        for attempt in range(self.max_retries + 1):
            try:
                result = self.agent_fn(
                    test_case.input,
                    timeout=self.timeout_seconds,
                )
                output = result.get("output", "")
                trajectory = result.get("trajectory", [])
                break
            except TimeoutError:
                error = f"Timeout after {self.timeout_seconds}s (attempt {attempt + 1})"
            except Exception as e:
                error = f"Error: {str(e)} (attempt {attempt + 1})"

        wall_time = time.time() - start_time

        return {
            "output": output,
            "trajectory": trajectory,
            "wall_time": wall_time,
            "tokens_used": result.get("tokens_used", 0) if not error else 0,
            "error": error,
        }

    def evaluate_single(self, test_case: TestCase) -> EvaluationResult:
        """Run and evaluate a single test case."""
        # Run the agent
        run_result = self.run_agent(test_case)

        # Apply all evaluators
        scores = {}
        for eval_name, eval_fn in self.evaluators.items():
            try:
                score = eval_fn(
                    input=test_case.input,
                    output=run_result["output"],
                    expected=test_case.expected_output,
                    trajectory=run_result["trajectory"],
                )
                scores[eval_name] = score
            except Exception as e:
                scores[eval_name] = 0.0
                print(f"  Evaluator '{eval_name}' failed: {e}")

        return EvaluationResult(
            test_case_id=test_case.id,
            agent_output=run_result["output"],
            scores=scores,
            trajectory=run_result["trajectory"],
            wall_time_seconds=run_result["wall_time"],
            tokens_used=run_result["tokens_used"],
            error=run_result["error"],
        )

    def run_evaluation(self, test_cases: list[TestCase]) -> dict:
        """Run the full evaluation pipeline.

        Args:
            test_cases: List of test cases to evaluate.

        Returns:
            Comprehensive evaluation report.
        """
        print(f"Starting evaluation: {len(test_cases)} test cases")
        print(f"Evaluators: {list(self.evaluators.keys())}")
        print("=" * 60)

        self.results = []
        for i, test_case in enumerate(test_cases):
            print(f"[{i+1}/{len(test_cases)}] Evaluating {test_case.id}...", end=" ")
            result = self.evaluate_single(test_case)
            self.results.append(result)

            if result.error:
                print(f"ERROR: {result.error}")
            else:
                score_str = ", ".join(
                    f"{k}: {v:.2f}" for k, v in result.scores.items()
                )
                print(f"Scores: {score_str} ({result.wall_time_seconds:.1f}s)")

        report = self.generate_report(test_cases)
        return report

    def generate_report(self, test_cases: list[TestCase]) -> dict:
        """Generate a comprehensive evaluation report."""
        if not self.results:
            return {"error": "No results to report."}

        # Overall metrics
        all_scores = {}
        for eval_name in self.evaluators:
            scores = [r.scores.get(eval_name, 0.0) for r in self.results]
            all_scores[eval_name] = {
                "mean": sum(scores) / len(scores),
                "min": min(scores),
                "max": max(scores),
                "std": (sum((s - sum(scores)/len(scores))**2 for s in scores) / len(scores)) ** 0.5,
            }

        # Error rate
        errors = [r for r in self.results if r.error]

        # Timing
        times = [r.wall_time_seconds for r in self.results]
        tokens = [r.tokens_used for r in self.results]

        # Per-category breakdown
        categories = {}
        for tc, result in zip(test_cases, self.results):
            cat = tc.category
            if cat not in categories:
                categories[cat] = []
            categories[cat].append(result)

        category_scores = {}
        for cat, cat_results in categories.items():
            for eval_name in self.evaluators:
                scores = [r.scores.get(eval_name, 0.0) for r in cat_results]
                key = f"{cat}/{eval_name}"
                category_scores[key] = sum(scores) / len(scores) if scores else 0.0

        # Per-difficulty breakdown
        difficulties = {}
        for tc, result in zip(test_cases, self.results):
            diff = tc.difficulty
            if diff not in difficulties:
                difficulties[diff] = []
            difficulties[diff].append(result)

        difficulty_scores = {}
        for diff, diff_results in difficulties.items():
            for eval_name in self.evaluators:
                scores = [r.scores.get(eval_name, 0.0) for r in diff_results]
                key = f"{diff}/{eval_name}"
                difficulty_scores[key] = sum(scores) / len(scores) if scores else 0.0

        report = {
            "timestamp": datetime.now().isoformat(),
            "total_test_cases": len(self.results),
            "overall_scores": all_scores,
            "error_count": len(errors),
            "error_rate": len(errors) / len(self.results),
            "timing": {
                "mean_seconds": sum(times) / len(times),
                "total_seconds": sum(times),
                "max_seconds": max(times),
            },
            "tokens": {
                "mean": sum(tokens) / len(tokens) if tokens else 0,
                "total": sum(tokens),
            },
            "by_category": category_scores,
            "by_difficulty": difficulty_scores,
        }

        return report

    def save_report(self, report: dict, path: str) -> None:
        """Save the evaluation report to a JSON file."""
        with open(path, "w") as f:
            json.dump(report, f, indent=2)
        print(f"Report saved to {path}")

ANNOTATION_GUIDELINES = """
## Agent Response Evaluation Guidelines

### Task
You will evaluate AI agent responses to user tasks. For each response,
rate the following dimensions on a 1-5 scale.

### Dimensions

**1. Task Completion (1-5)**
- 5: Fully completes the task with no omissions
- 4: Mostly complete with minor omissions
- 3: Partially complete; addresses the main point but misses aspects
- 2: Minimally addresses the task
- 1: Does not address the task at all

**2. Accuracy (1-5)**
- 5: All information is accurate
- 4: Nearly all accurate, one minor error
- 3: Mostly accurate, some errors
- 2: Multiple significant errors
- 1: Predominantly inaccurate

**3. Helpfulness (1-5)**
- 5: Exceptionally helpful; would fully satisfy the user
- 4: Helpful; addresses the user's needs well
- 3: Somewhat helpful; provides partial value
- 2: Minimally helpful
- 1: Not helpful at all

### Instructions
- Read the task description carefully before evaluating.
- Evaluate the response independently (do not compare to other responses).
- Provide a brief justification (1-2 sentences) for each score.
- If you are unsure, err on the side of a lower score.
- Flag any responses that contain harmful, biased, or inappropriate content.
"""

def cohens_kappa(rater1: list[int], rater2: list[int]) -> float:
    """Compute Cohen's Kappa for inter-rater reliability.

    Cohen's Kappa measures agreement between two raters beyond
    what would be expected by chance.

    Interpretation:
    - < 0.00: Poor agreement
    - 0.00-0.20: Slight agreement
    - 0.21-0.40: Fair agreement
    - 0.41-0.60: Moderate agreement
    - 0.61-0.80: Substantial agreement
    - 0.81-1.00: Almost perfect agreement

    Args:
        rater1: List of ratings from rater 1.
        rater2: List of ratings from rater 2.

    Returns:
        Cohen's Kappa coefficient.
    """
    assert len(rater1) == len(rater2), "Raters must evaluate the same items"

    n = len(rater1)
    categories = sorted(set(rater1) | set(rater2))

    # Observed agreement
    agreement = sum(1 for a, b in zip(rater1, rater2) if a == b)
    p_observed = agreement / n

    # Expected agreement by chance
    p_expected = 0.0
    for cat in categories:
        p1 = sum(1 for r in rater1 if r == cat) / n
        p2 = sum(1 for r in rater2 if r == cat) / n
        p_expected += p1 * p2

    # Kappa
    if p_expected == 1.0:
        return 1.0  # Perfect agreement
    return (p_observed - p_expected) / (1 - p_expected)


def krippendorffs_alpha(ratings: list[list[int | None]], level: str = "ordinal") -> float:
    """Compute Krippendorff's Alpha for multiple raters.

    Krippendorff's Alpha supports:
    - Any number of raters
    - Missing data (None values)
    - Different measurement levels (nominal, ordinal, interval, ratio)

    This is a simplified implementation for nominal data.

    Args:
        ratings: Matrix where ratings[i][j] is rater i's rating for item j.
                 None indicates a missing rating.
        level: Measurement level ("nominal" supported in this implementation).

    Returns:
        Krippendorff's Alpha coefficient.
    """
    n_raters = len(ratings)
    n_items = len(ratings[0])

    # Collect observed disagreement
    observed_disagreement = 0.0
    n_pairs = 0

    for j in range(n_items):
        item_ratings = [ratings[i][j] for i in range(n_raters) if ratings[i][j] is not None]
        m = len(item_ratings)
        if m < 2:
            continue
        for a_idx in range(m):
            for b_idx in range(a_idx + 1, m):
                if item_ratings[a_idx] != item_ratings[b_idx]:
                    observed_disagreement += 1
                n_pairs += 1

    if n_pairs == 0:
        return 1.0

    d_observed = observed_disagreement / n_pairs

    # Expected disagreement
    all_ratings = [r for rater in ratings for item_ratings in zip(*ratings) for r in [rater[ratings[0].index(item_ratings[0])]] if r is not None]
    all_valid = [r for rater in ratings for r in rater if r is not None]
    n_total = len(all_valid)
    value_counts = {}
    for v in all_valid:
        value_counts[v] = value_counts.get(v, 0) + 1

    d_expected = 0.0
    values = list(value_counts.keys())
    for i, v1 in enumerate(values):
        for j, v2 in enumerate(values):
            if v1 != v2:
                d_expected += value_counts[v1] * value_counts[v2]
    d_expected /= (n_total * (n_total - 1))

    if d_expected == 0:
        return 1.0

    return 1.0 - (d_observed / d_expected)

# Approximate pricing (as of early 2025)
MODEL_PRICING = {
    "gpt-4-turbo": {"input": 10.00, "output": 30.00},     # per million tokens
    "gpt-4o": {"input": 2.50, "output": 10.00},
    "gpt-4o-mini": {"input": 0.15, "output": 0.60},
    "claude-3-5-sonnet": {"input": 3.00, "output": 15.00},
    "claude-3-5-haiku": {"input": 0.80, "output": 4.00},
}


@dataclass
class CostMetrics:
    """Track costs for an agent interaction."""

    model: str
    input_tokens: int = 0
    output_tokens: int = 0
    tool_calls: int = 0
    tool_cost: float = 0.0  # External API costs

    @property
    def llm_cost(self) -> float:
        """Compute LLM API cost in USD."""
        pricing = MODEL_PRICING.get(self.model, {"input": 0, "output": 0})
        input_cost = (self.input_tokens / 1_000_000) * pricing["input"]
        output_cost = (self.output_tokens / 1_000_000) * pricing["output"]
        return input_cost + output_cost

    @property
    def total_cost(self) -> float:
        return self.llm_cost + self.tool_cost

    def __repr__(self) -> str:
        return (
            f"CostMetrics(model={self.model}, "
            f"tokens={self.input_tokens}+{self.output_tokens}, "
            f"llm_cost=${self.llm_cost:.4f}, "
            f"total=${self.total_cost:.4f})"
        )


def cost_adjusted_score(
    task_score: float,
    cost: float,
    max_acceptable_cost: float = 1.0,
    cost_penalty_weight: float = 0.2,
) -> float:
    """Compute a cost-adjusted performance score.

    Penalizes agents that are expensive relative to their accuracy.

    Args:
        task_score: Raw task performance score (0-1).
        cost: Cost of the interaction in USD.
        max_acceptable_cost: Cost threshold above which penalty increases.
        cost_penalty_weight: How much to weight cost penalty (0-1).

    Returns:
        Adjusted score.
    """
    cost_ratio = min(cost / max_acceptable_cost, 2.0)  # Cap at 2x
    cost_penalty = cost_ratio * cost_penalty_weight
    return max(0.0, task_score - cost_penalty)

@dataclass
class LatencyMetrics:
    """Track latency for an agent interaction."""

    time_to_first_token: float = 0.0     # Seconds until first output
    time_to_completion: float = 0.0       # Total interaction time
    llm_call_times: list[float] = field(default_factory=list)
    tool_call_times: list[float] = field(default_factory=list)

    @property
    def total_llm_time(self) -> float:
        return sum(self.llm_call_times)

    @property
    def total_tool_time(self) -> float:
        return sum(self.tool_call_times)

    @property
    def overhead_time(self) -> float:
        """Time spent on orchestration (not LLM or tools)."""
        return self.time_to_completion - self.total_llm_time - self.total_tool_time

    def summary(self) -> dict:
        return {
            "total_seconds": self.time_to_completion,
            "time_to_first_token": self.time_to_first_token,
            "llm_time": self.total_llm_time,
            "tool_time": self.total_tool_time,
            "overhead_time": self.overhead_time,
            "num_llm_calls": len(self.llm_call_times),
            "num_tool_calls": len(self.tool_call_times),
            "avg_llm_call_time": (
                self.total_llm_time / len(self.llm_call_times)
                if self.llm_call_times else 0
            ),
        }

def find_pareto_frontier(
    agents: list[dict],
) -> list[dict]:
    """Find the Pareto-optimal agents (cost vs. quality).

    An agent is Pareto-optimal if no other agent is both cheaper
    and higher quality.

    Args:
        agents: List of dicts with 'name', 'quality', and 'cost' keys.

    Returns:
        List of Pareto-optimal agents.
    """
    # Sort by cost (ascending)
    sorted_agents = sorted(agents, key=lambda a: a["cost"])

    pareto = []
    best_quality = -1

    for agent in sorted_agents:
        if agent["quality"] > best_quality:
            pareto.append(agent)
            best_quality = agent["quality"]

    return pareto

"""
A complete automated evaluation harness for AI agents.

This harness:
1. Loads test cases from a JSON file
2. Runs the agent on each test case
3. Applies multiple evaluation metrics
4. Generates a detailed report with breakdowns

Requirements:
    pip install sentence-transformers
"""

import json
import time
from dataclasses import dataclass, field
from datetime import datetime


# ── Evaluator Functions ────────────────────────────────────────────

def exact_match_evaluator(
    input: str, output: str, expected: str | None, **kwargs
) -> float:
    """Check if output exactly matches expected output.

    Case-insensitive, whitespace-normalized comparison.
    """
    if expected is None:
        return 0.0
    return 1.0 if output.strip().lower() == expected.strip().lower() else 0.0


def contains_evaluator(
    input: str, output: str, expected: str | None, **kwargs
) -> float:
    """Check if the expected answer is contained in the output."""
    if expected is None:
        return 0.0
    return 1.0 if expected.strip().lower() in output.strip().lower() else 0.0


def semantic_similarity_evaluator(
    input: str, output: str, expected: str | None, **kwargs
) -> float:
    """Compute semantic similarity between output and expected.

    Uses sentence-transformers for embedding-based comparison.
    """
    if expected is None:
        return 0.0

    from sentence_transformers import SentenceTransformer
    import numpy as np

    # Use a cached model in production
    model = SentenceTransformer("all-MiniLM-L6-v2")
    embeddings = model.encode([output, expected])
    similarity = float(np.dot(embeddings[0], embeddings[1]) / (
        np.linalg.norm(embeddings[0]) * np.linalg.norm(embeddings[1]) + 1e-8
    ))
    return max(0.0, similarity)  # Clamp negative similarities


def step_efficiency_evaluator(
    input: str, output: str, expected: str | None,
    trajectory: list[dict] | None = None, **kwargs
) -> float:
    """Evaluate how efficiently the agent reached its answer.

    Penalizes excessive steps, errors, and backtracking.
    """
    if not trajectory:
        return 0.5  # Unknown efficiency

    n_steps = len(trajectory)
    n_errors = sum(1 for step in trajectory if step.get("status") == "error")
    n_retries = sum(
        1 for i in range(1, len(trajectory))
        if trajectory[i].get("action") == trajectory[i-1].get("action")
    )

    # Score: penalize for excessive steps, errors, and retries
    step_penalty = max(0, n_steps - 5) * 0.05  # Penalty after 5 steps
    error_penalty = n_errors * 0.15
    retry_penalty = n_retries * 0.10

    score = max(0.0, 1.0 - step_penalty - error_penalty - retry_penalty)
    return score


# ── Evaluation Harness ─────────────────────────────────────────────

class AgentEvaluationHarness:
    """A comprehensive evaluation harness for AI agents.

    Features:
    - Multiple evaluator support
    - Detailed per-test-case results
    - Aggregate statistics with breakdowns by category and difficulty
    - Reproducible evaluation with seeded runs
    """

    def __init__(self, agent_fn):
        self.agent_fn = agent_fn
        self.evaluators = {
            "exact_match": exact_match_evaluator,
            "contains": contains_evaluator,
            "efficiency": step_efficiency_evaluator,
        }
        self.results = []

    def add_evaluator(self, name: str, evaluator_fn) -> None:
        """Register a custom evaluator."""
        self.evaluators[name] = evaluator_fn

    def run(
        self,
        test_cases: list[dict],
        num_runs: int = 1,
        verbose: bool = True,
    ) -> dict:
        """Run the evaluation harness.

        Args:
            test_cases: List of test case dicts.
            num_runs: Number of times to run each test case (for variance).
            verbose: Whether to print progress.

        Returns:
            Comprehensive evaluation report.
        """
        if verbose:
            print(f"Evaluation started: {len(test_cases)} cases, {num_runs} run(s) each")
            print(f"Evaluators: {list(self.evaluators.keys())}")
            print("=" * 70)

        self.results = []

        for i, test_case in enumerate(test_cases):
            case_id = test_case.get("id", f"case_{i}")
            if verbose:
                print(f"\n[{i+1}/{len(test_cases)}] {case_id}: {test_case['input'][:60]}...")

            run_results = []
            for run in range(num_runs):
                # Run the agent
                start = time.time()
                try:
                    agent_result = self.agent_fn(test_case["input"])
                    output = agent_result if isinstance(agent_result, str) else agent_result.get("output", "")
                    trajectory = agent_result.get("trajectory", []) if isinstance(agent_result, dict) else []
                    error = None
                except Exception as e:
                    output = ""
                    trajectory = []
                    error = str(e)
                elapsed = time.time() - start

                # Apply evaluators
                scores = {}
                for eval_name, eval_fn in self.evaluators.items():
                    try:
                        score = eval_fn(
                            input=test_case["input"],
                            output=output,
                            expected=test_case.get("expected_output"),
                            trajectory=trajectory,
                        )
                        scores[eval_name] = score
                    except Exception as e:
                        scores[eval_name] = 0.0

                run_results.append({
                    "output": output,
                    "scores": scores,
                    "wall_time": elapsed,
                    "error": error,
                })

            # Aggregate across runs
            avg_scores = {}
            for eval_name in self.evaluators:
                run_scores = [r["scores"].get(eval_name, 0.0) for r in run_results]
                avg_scores[eval_name] = {
                    "mean": sum(run_scores) / len(run_scores),
                    "std": (sum((s - sum(run_scores)/len(run_scores))**2 for s in run_scores) / len(run_scores)) ** 0.5 if num_runs > 1 else 0.0,
                }

            case_result = {
                "test_case_id": case_id,
                "category": test_case.get("category", "general"),
                "difficulty": test_case.get("difficulty", "medium"),
                "scores": avg_scores,
                "num_runs": num_runs,
                "avg_wall_time": sum(r["wall_time"] for r in run_results) / num_runs,
                "error_rate": sum(1 for r in run_results if r["error"]) / num_runs,
                "sample_output": run_results[0]["output"][:300],
            }
            self.results.append(case_result)

            if verbose:
                score_str = " | ".join(
                    f"{k}: {v['mean']:.2f}" for k, v in avg_scores.items()
                )
                print(f"  Scores: {score_str} ({case_result['avg_wall_time']:.1f}s)")

        # Generate report
        report = self._generate_report()

        if verbose:
            print("\n" + "=" * 70)
            print("EVALUATION SUMMARY")
            print("=" * 70)
            self._print_report(report)

        return report

    def _generate_report(self) -> dict:
        """Generate the evaluation report."""
        # Overall scores
        overall = {}
        for eval_name in self.evaluators:
            means = [r["scores"][eval_name]["mean"] for r in self.results]
            overall[eval_name] = {
                "mean": sum(means) / len(means),
                "std": (sum((m - sum(means)/len(means))**2 for m in means) / len(means)) ** 0.5,
                "min": min(means),
                "max": max(means),
            }

        # By category
        by_category = {}
        for result in self.results:
            cat = result["category"]
            if cat not in by_category:
                by_category[cat] = []
            by_category[cat].append(result)

        category_scores = {}
        for cat, cat_results in by_category.items():
            category_scores[cat] = {}
            for eval_name in self.evaluators:
                means = [r["scores"][eval_name]["mean"] for r in cat_results]
                category_scores[cat][eval_name] = sum(means) / len(means)
            category_scores[cat]["count"] = len(cat_results)

        # By difficulty
        by_difficulty = {}
        for result in self.results:
            diff = result["difficulty"]
            if diff not in by_difficulty:
                by_difficulty[diff] = []
            by_difficulty[diff].append(result)

        difficulty_scores = {}
        for diff, diff_results in by_difficulty.items():
            difficulty_scores[diff] = {}
            for eval_name in self.evaluators:
                means = [r["scores"][eval_name]["mean"] for r in diff_results]
                difficulty_scores[diff][eval_name] = sum(means) / len(means)
            difficulty_scores[diff]["count"] = len(diff_results)

        # Timing
        times = [r["avg_wall_time"] for r in self.results]

        return {
            "timestamp": datetime.now().isoformat(),
            "total_cases": len(self.results),
            "overall_scores": overall,
            "by_category": category_scores,
            "by_difficulty": difficulty_scores,
            "timing": {
                "mean_seconds": sum(times) / len(times),
                "total_seconds": sum(times),
                "max_seconds": max(times),
            },
            "error_rate": sum(r["error_rate"] for r in self.results) / len(self.results),
        }

    def _print_report(self, report: dict) -> None:
        """Pretty-print the evaluation report."""
        print(f"\nTotal test cases: {report['total_cases']}")
        print(f"Error rate: {report['error_rate']:.1%}")
        print(f"Mean time per case: {report['timing']['mean_seconds']:.1f}s")

        print("\nOverall Scores:")
        for metric, stats in report["overall_scores"].items():
            print(f"  {metric}: {stats['mean']:.3f} (std: {stats['std']:.3f}, range: {stats['min']:.2f}-{stats['max']:.2f})")

        if report["by_category"]:
            print("\nBy Category:")
            for cat, scores in report["by_category"].items():
                count = scores.pop("count", "?")
                score_str = ", ".join(f"{k}: {v:.2f}" for k, v in scores.items())
                print(f"  {cat} (n={count}): {score_str}")

        if report["by_difficulty"]:
            print("\nBy Difficulty:")
            for diff, scores in report["by_difficulty"].items():
                count = scores.pop("count", "?")
                score_str = ", ".join(f"{k}: {v:.2f}" for k, v in scores.items())
                print(f"  {diff} (n={count}): {score_str}")


# ── Usage Example ─────────────────────────────────────────────────

def main():
    """Demonstrate the evaluation harness."""

    # A simple mock agent for demonstration
    def mock_agent(input_text: str) -> dict:
        """A mock agent that returns predefined answers."""
        answers = {
            "capital of france": "The capital of France is Paris.",
            "2 + 2": "4",
            "sort": "def sort_list(lst): return sorted(lst)",
        }

        output = "I don't know."
        for key, answer in answers.items():
            if key in input_text.lower():
                output = answer
                break

        return {
            "output": output,
            "trajectory": [
                {"action": "think", "content": "Processing query..."},
                {"action": "respond", "content": output},
            ],
        }

    # Define test cases
    test_cases = [
        {
            "id": "geo_001",
            "input": "What is the capital of France?",
            "expected_output": "Paris",
            "category": "geography",
            "difficulty": "easy",
        },
        {
            "id": "math_001",
            "input": "What is 2 + 2?",
            "expected_output": "4",
            "category": "math",
            "difficulty": "easy",
        },
        {
            "id": "code_001",
            "input": "Write a Python function to sort a list.",
            "expected_output": "def sort_list(lst): return sorted(lst)",
            "category": "coding",
            "difficulty": "medium",
        },
        {
            "id": "geo_002",
            "input": "What is the capital of Bhutan?",
            "expected_output": "Thimphu",
            "category": "geography",
            "difficulty": "hard",
        },
        {
            "id": "reason_001",
            "input": "If all roses are flowers and some flowers fade quickly, can we conclude that some roses fade quickly?",
            "expected_output": "No",
            "category": "reasoning",
            "difficulty": "hard",
        },
    ]

    # Create and run harness
    harness = AgentEvaluationHarness(mock_agent)
    report = harness.run(test_cases, num_runs=1, verbose=True)

    # Save report
    print(f"\nFull report:")
    print(json.dumps(report, indent=2, default=str))


if __name__ == "__main__":
    main()

Evaluation started: 5 cases, 1 run(s) each
Evaluators: ['exact_match', 'contains', 'efficiency']
======================================================================

[1/5] geo_001: What is the capital of France?...
  Scores: exact_match: 0.00 | contains: 1.00 | efficiency: 1.00 (0.0s)

[2/5] math_001: What is 2 + 2?...
  Scores: exact_match: 1.00 | contains: 1.00 | efficiency: 1.00 (0.0s)

[3/5] code_001: Write a Python function to sort a list....
  Scores: exact_match: 1.00 | contains: 1.00 | efficiency: 1.00 (0.0s)

[4/5] geo_002: What is the capital of Bhutan?...
  Scores: exact_match: 0.00 | contains: 0.00 | efficiency: 1.00 (0.0s)

[5/5] reason_001: If all roses are flowers and some flowers fade quickly, ...
  Scores: exact_match: 0.00 | contains: 0.00 | efficiency: 1.00 (0.0s)

======================================================================
EVALUATION SUMMARY
======================================================================

Total test cases: 5
Error rate: 0.0%
Mean time per case: 0.0s

Overall Scores:
  exact_match: 0.400 (std: 0.490, range: 0.00-1.00)
  contains: 0.600 (std: 0.490, range: 0.00-1.00)
  efficiency: 1.000 (std: 0.000, range: 1.00-1.00)

By Category:
  geography (n=2): exact_match: 0.00, contains: 0.50
  math (n=1): exact_match: 1.00, contains: 1.00
  coding (n=1): exact_match: 1.00, contains: 1.00
  reasoning (n=1): exact_match: 0.00, contains: 0.00

By Difficulty:
  easy (n=2): exact_match: 0.50, contains: 1.00
  medium (n=1): exact_match: 1.00, contains: 1.00
  hard (n=2): exact_match: 0.00, contains: 0.00