import asyncio
from dataclasses import dataclass
 
@dataclass
class CandidateSolution:
    code_diff: str
    test_results: dict
    score: float
 
async def best_of_n(task: str, n: int, agent_fn, score_fn) -> CandidateSolution:
    """Generate N candidate solutions and return the best one."""
    # Launch all candidates in parallel
    candidates = await asyncio.gather(
        *[agent_fn(task) for _ in range(n)]
    )
 
    # Score each candidate
    scored = []
    for candidate in candidates:
        score = await score_fn(candidate)
        scored.append(CandidateSolution(
            code_diff=candidate.diff,
            test_results=candidate.tests,
            score=score,
        ))
 
    # Return the highest-scoring candidate
    return max(scored, key=lambda c: c.score)

P(at least one success in k attempts) = 1 - (1 - p)^k

async def hybrid_best_of_n(task: str, n: int, top_k: int, agent_fn, refine_fn, score_fn):
    """Best-of-N with a refinement phase for top candidates."""
    # Phase 1: Generate N diverse candidates
    candidates = await asyncio.gather(
        *[agent_fn(task) for _ in range(n)]
    )
    scored = [(c, await score_fn(c)) for c in candidates]
    scored.sort(key=lambda x: x[1], reverse=True)
 
    # Phase 2: Refine top-K candidates
    top_candidates = [c for c, _ in scored[:top_k]]
    refined = await asyncio.gather(
        *[refine_fn(c, task) for c in top_candidates]
    )
 
    # Phase 3: Final scoring
    final_scored = [(r, await score_fn(r)) for r in refined]
    return max(final_scored, key=lambda x: x[1])

# docker-compose.yml for Best-of-N candidate environments
version: "3.8"
services:
  candidate-template:
    build: .
    volumes:
      - ./repo:/workspace:ro  # Read-only source mount
    tmpfs:
      - /workspace-writable:size=2G
    environment:
      - CANDIDATE_ID=${CANDIDATE_ID}
    command: >
      bash -c "
        cp -r /workspace/* /workspace-writable/ &&
        cd /workspace-writable &&
        exec agent-runner --task '${TASK_DESCRIPTION}'
      "

import asyncio
import logging
import time
from typing import Optional
 
logger = logging.getLogger(__name__)
 
class BestOfNOrchestrator:
    def __init__(
        self,
        n: int = 5,
        timeout_seconds: int = 600,
        max_concurrent: int = 5,
    ):
        self.n = n
        self.timeout = timeout_seconds
        self.max_concurrent = max_concurrent
        self.semaphore = asyncio.Semaphore(max_concurrent)
 
    async def _run_candidate(self, task: str, candidate_id: int, agent_fn):
        """Run a single candidate with timeout and error handling."""
        async with self.semaphore:
            start = time.time()
            try:
                result = await asyncio.wait_for(
                    agent_fn(task, candidate_id=candidate_id),
                    timeout=self.timeout,
                )
                elapsed = time.time() - start
                logger.info(f"Candidate {candidate_id} completed in {elapsed:.1f}s")
                return result
            except asyncio.TimeoutError:
                logger.warning(f"Candidate {candidate_id} timed out")
                return None
            except Exception as e:
                logger.error(f"Candidate {candidate_id} failed: {e}")
                return None
 
    async def run(self, task: str, agent_fn, score_fn) -> Optional[dict]:
        """Execute Best-of-N and return the best candidate."""
        logger.info(f"Starting Best-of-N with N={self.n} for task: {task[:80]}...")
 
        # Launch all candidates
        results = await asyncio.gather(
            *[
                self._run_candidate(task, i, agent_fn)
                for i in range(self.n)
            ]
        )
 
        # Filter out failures
        valid = [r for r in results if r is not None]
        if not valid:
            logger.error("All candidates failed")
            return None
 
        logger.info(f"{len(valid)}/{self.n} candidates completed successfully")
 
        # Score and rank
        scored = []
        for result in valid:
            score = await score_fn(result)
            scored.append({"result": result, "score": score})
 
        scored.sort(key=lambda x: x["score"], reverse=True)
 
        winner = scored[0]
        logger.info(f"Best candidate score: {winner['score']:.3f}")
        return winner["result"]

async def diverse_agent_fn(task: str, candidate_id: int):
    """Agent function with temperature variation for diversity."""
    temperatures = [0.3, 0.5, 0.7, 0.9, 1.0]
    temp = temperatures[candidate_id % len(temperatures)]
 
    return await call_agent(
        task=task,
        temperature=temp,
        system_prompt=BASE_SYSTEM_PROMPT,
    )

APPROACH_HINTS = [
    "Try to solve this with minimal changes to existing code.",
    "Consider whether a refactoring approach might be cleaner here.",
    "Focus on adding proper error handling as part of your fix.",
    "Look for the root cause rather than patching the symptom.",
    "Consider edge cases that the current tests might not cover.",
]
 
async def diverse_prompt_agent(task: str, candidate_id: int):
    hint = APPROACH_HINTS[candidate_id % len(APPROACH_HINTS)]
    augmented_task = f"{task}\n\nApproach hint: {hint}"
    return await call_agent(task=augmented_task)

async def test_based_score(candidate) -> float:
    """Score based on test pass rate."""
    test_result = await run_test_suite(candidate.workspace)
 
    if test_result.all_passed:
        return 1.0
 
    pass_rate = test_result.passed / test_result.total
    return pass_rate

async def multi_signal_score(candidate) -> float:
    """Composite score from multiple quality signals."""
    scores = {}
 
    # 1. Test pass rate (weighted heavily)
    test_result = await run_test_suite(candidate.workspace)
    scores["tests"] = test_result.passed / max(test_result.total, 1)
 
    # 2. Diff size (prefer smaller, focused changes)
    diff_lines = len(candidate.diff.splitlines())
    scores["diff_size"] = max(0, 1.0 - (diff_lines / 500))
 
    # 3. Linter compliance
    lint_result = await run_linter(candidate.workspace)
    lint_issues = lint_result.error_count + lint_result.warning_count
    scores["lint"] = max(0, 1.0 - (lint_issues / 20))
 
    # 4. Type checker
    type_result = await run_type_checker(candidate.workspace)
    scores["types"] = 1.0 if type_result.success else 0.5
 
    # 5. No regressions (existing tests still pass)
    regression_result = await run_existing_tests(candidate.workspace)
    scores["regression"] = 1.0 if regression_result.all_passed else 0.0
 
    # Weighted combination
    weights = {
        "tests": 0.35,
        "diff_size": 0.10,
        "lint": 0.10,
        "types": 0.10,
        "regression": 0.35,
    }
 
    final = sum(scores[k] * weights[k] for k in weights)
    return final

async def llm_judge_score(candidate, task: str) -> float:
    """Use an LLM to evaluate code quality aspects."""
    prompt = f"""You are evaluating a code change made by an AI coding agent.
 
Task description:
{task}
 
Code diff:
{candidate.diff}
 
Rate this solution on a scale of 1-10 for each criterion:
1. Correctness: Does it actually solve the stated problem?
2. Minimal change: Does it avoid unnecessary modifications?
3. Readability: Is the code clear and well-structured?
4. Robustness: Does it handle edge cases?
5. Consistency: Does it match the existing code style?
 
Respond with JSON: {{"correctness": N, "minimal": N, "readability": N, "robustness": N, "consistency": N}}"""
 
    response = await call_llm(prompt, model="claude-sonnet-4", temperature=0)
    scores = parse_json(response)
 
    # Normalize to 0-1
    return sum(scores.values()) / (5 * 10)

from difflib import SequenceMatcher
 
def find_consensus(candidates: list) -> Optional[dict]:
    """Find the candidate most similar to all others."""
    if len(candidates) <= 1:
        return candidates[0] if candidates else None
 
    # Compare each candidate's diff to all others
    similarity_scores = []
    for i, c1 in enumerate(candidates):
        total_similarity = sum(
            SequenceMatcher(None, c1.diff, c2.diff).ratio()
            for j, c2 in enumerate(candidates) if i != j
        )
        avg_similarity = total_similarity / (len(candidates) - 1)
        similarity_scores.append((c1, avg_similarity))
 
    # Return the candidate closest to the group consensus
    return max(similarity_scores, key=lambda x: x[1])[0]

import asyncio
from openai import AsyncOpenAI
 
client = AsyncOpenAI()
 
async def codex_best_of_n(
    task_description: str,
    repo_url: str,
    n: int = 5,
) -> dict:
    """Run N parallel Codex tasks and select the best result."""
 
    async def launch_single_task(candidate_id: int):
        """Launch one Codex coding task."""
        response = await client.responses.create(
            model="codex-mini-latest",
            instructions=(
                f"You are solving a coding task. Candidate run #{candidate_id}.\n"
                f"Repository: {repo_url}\n\n"
                f"Task: {task_description}\n\n"
                "After implementing your solution, run the test suite to verify."
            ),
            tools=[
                {"type": "code_interpreter"},
            ],
            temperature=0.4 + (candidate_id * 0.15),  # Vary temperature
        )
        return {
            "candidate_id": candidate_id,
            "response": response,
            "output": response.output_text,
        }
 
    # Launch all candidates in parallel
    results = await asyncio.gather(
        *[launch_single_task(i) for i in range(n)],
        return_exceptions=True,
    )
 
    # Filter successful results
    valid = [r for r in results if isinstance(r, dict)]
 
    if not valid:
        raise RuntimeError("All Codex candidates failed")
 
    # Score and select (simplified: prefer candidates that mention "tests pass")
    for r in valid:
        output_lower = r["output"].lower()
        r["score"] = 0.0
        if "all tests pass" in output_lower or "tests passed" in output_lower:
            r["score"] += 1.0
        if "error" in output_lower or "failed" in output_lower:
            r["score"] -= 0.5
 
    valid.sort(key=lambda r: r["score"], reverse=True)
    return valid[0]

async def best_of_n_pr_workflow(issue_url: str, n: int = 3):
    """Generate N candidate PRs and pick the best one."""
    issue = await fetch_github_issue(issue_url)
 
    candidates = await asyncio.gather(
        *[
            codex_create_pr_candidate(issue, candidate_id=i)
            for i in range(n)
        ]
    )
 
    # Score each candidate PR
    for candidate in candidates:
        candidate["score"] = await score_pr_candidate(candidate)
 
    # Select the best
    best = max(candidates, key=lambda c: c["score"])
 
    # Create the actual PR from the best candidate
    pr_url = await create_github_pr(
        branch=best["branch"],
        title=best["title"],
        body=best["body"],
    )
 
    return pr_url

# CLAUDE.md configuration for Best-of-N workflow
## Best-of-N Instructions
 
When tackling complex bug fixes or feature implementations:
1. First, analyze the problem and identify 2-3 possible approaches
2. For each approach, create a separate worktree branch
3. Implement each approach independently
4. Run the test suite against each approach
5. Compare results and select the best implementation
 
## Approach Diversity
- Approach A: Minimal targeted fix
- Approach B: Refactoring-based solution
- Approach C: Defensive coding with additional validation

import asyncio
import anthropic
 
client = anthropic.AsyncAnthropic()
 
async def claude_best_of_n(
    task: str,
    codebase_context: str,
    n: int = 5,
) -> dict:
    """Run Best-of-N using Claude for code generation."""
 
    async def generate_candidate(candidate_id: int):
        temp = 0.3 + (candidate_id * 0.1)
        response = await client.messages.create(
            model="claude-sonnet-4-20250514",
            max_tokens=8096,
            temperature=min(temp, 1.0),
            system=(
                "You are an expert software engineer. Generate a complete, "
                "working solution. Output only the code diff in unified diff format."
            ),
            messages=[
                {
                    "role": "user",
                    "content": (
                        f"Codebase context:\n{codebase_context}\n\n"
                        f"Task:\n{task}\n\n"
                        "Provide your solution as a unified diff."
                    ),
                }
            ],
        )
        return {
            "candidate_id": candidate_id,
            "diff": response.content[0].text,
            "temperature": temp,
        }
 
    candidates = await asyncio.gather(
        *[generate_candidate(i) for i in range(n)]
    )
 
    # Apply each diff in isolation and test
    scored = []
    for candidate in candidates:
        test_result = await apply_and_test(candidate["diff"])
        candidate["test_passed"] = test_result.passed
        candidate["test_count"] = test_result.total
        candidate["score"] = test_result.passed / max(test_result.total, 1)
        scored.append(candidate)
 
    scored.sort(key=lambda c: c["score"], reverse=True)
    return scored[0]

#!/bin/bash
# best_of_n_swe_agent.sh - Run SWE-agent N times and pick the best
 
TASK_INSTANCE="django__django-16379"
N=5
RESULTS_DIR="./results/${TASK_INSTANCE}"
mkdir -p "$RESULTS_DIR"
 
# Run N candidates in parallel
for i in $(seq 1 $N); do
    python run_swe_agent.py \
        --instance "$TASK_INSTANCE" \
        --seed "$i" \
        --output_dir "${RESULTS_DIR}/candidate_${i}" \
        --temperature 0.$(( 3 + i )) &
done
 
# Wait for all candidates to complete
wait
 
# Score each candidate
python score_candidates.py \
    --results_dir "$RESULTS_DIR" \
    --scoring_mode "test_pass_rate" \
    --output "${RESULTS_DIR}/best_candidate.json"

# openhands_best_of_n.py
from openhands.core.config import AppConfig
from openhands.core.main import run_agent
 
async def openhands_best_of_n(task: str, n: int = 5):
    """Run OpenHands agent N times with different configurations."""
    configs = []
    for i in range(n):
        config = AppConfig()
        config.llm.temperature = 0.3 + (i * 0.15)
        config.llm.top_p = 0.9 + (i * 0.02)
        config.sandbox.timeout = 300
        configs.append(config)
 
    results = await asyncio.gather(
        *[run_agent(task, config=c) for c in configs]
    )
 
    # Evaluate and return best
    scored = [(r, evaluate_result(r)) for r in results]
    return max(scored, key=lambda x: x[1])[0]

import asyncio
import subprocess
import json
import tempfile
import shutil
from pathlib import Path
 
class DockerBestOfN:
    def __init__(self, repo_path: str, n: int = 5):
        self.repo_path = Path(repo_path)
        self.n = n
        self.workdirs = []
 
    async def setup_workspaces(self):
        """Create N isolated copies of the repository."""
        for i in range(self.n):
            workdir = Path(tempfile.mkdtemp(prefix=f"bon_candidate_{i}_"))
            shutil.copytree(self.repo_path, workdir / "repo")
            self.workdirs.append(workdir / "repo")
 
    async def run_candidate(self, idx: int, task: str) -> dict:
        """Run a coding agent in an isolated workspace."""
        workdir = self.workdirs[idx]
 
        # Run agent in Docker container
        result = subprocess.run(
            [
                "docker", "run", "--rm",
                "-v", f"{workdir}:/workspace",
                "-e", f"TASK={task}",
                "-e", f"TEMPERATURE={0.3 + idx * 0.15}",
                "coding-agent:latest",
            ],
            capture_output=True,
            text=True,
            timeout=600,
        )
 
        # Get the diff
        diff_result = subprocess.run(
            ["git", "diff"],
            cwd=workdir,
            capture_output=True,
            text=True,
        )
 
        # Run tests
        test_result = subprocess.run(
            ["python", "-m", "pytest", "--tb=short", "-q"],
            cwd=workdir,
            capture_output=True,
            text=True,
        )
 
        return {
            "candidate_id": idx,
            "diff": diff_result.stdout,
            "test_output": test_result.stdout,
            "test_returncode": test_result.returncode,
            "agent_output": result.stdout,
        }
 
    async def run(self, task: str) -> dict:
        """Execute the full Best-of-N pipeline."""
        await self.setup_workspaces()
 
        results = await asyncio.gather(
            *[self.run_candidate(i, task) for i in range(self.n)]
        )
 
        # Score: tests passing is the primary signal
        for r in results:
            r["score"] = 1.0 if r["test_returncode"] == 0 else 0.0
            # Bonus for smaller diffs
            r["score"] += max(0, 0.2 - len(r["diff"]) / 10000)
 
        results.sort(key=lambda r: r["score"], reverse=True)
        return results[0]
 
    def cleanup(self):
        """Remove temporary workspaces."""
        for workdir in self.workdirs:
            shutil.rmtree(workdir.parent, ignore_errors=True)

async def adaptive_best_of_n(task: str, agent_fn, score_fn):
    """Adjust N based on estimated task difficulty."""
    difficulty = await estimate_difficulty(task)
 
    if difficulty == "easy":
        n = 1  # Single attempt is fine
    elif difficulty == "medium":
        n = 3
    elif difficulty == "hard":
        n = 5
    else:  # "very hard"
        n = 10
 
    return await best_of_n(task, n, agent_fn, score_fn)

async def early_stop_best_of_n(task: str, n: int, agent_fn, score_fn):
    """Stop early if a perfect candidate is found."""
    best_result = None
    best_score = -1
 
    for i in range(n):
        result = await agent_fn(task, candidate_id=i)
        score = await score_fn(result)
 
        if score > best_score:
            best_score = score
            best_result = result
 
        # Early termination on perfect score
        if score >= 1.0:
            break
 
    return best_result

async def cascading_best_of_n(task: str, agent_fn, score_fn, threshold: float = 0.9):
    """Start with 1 attempt, escalate to N if needed."""
    # First attempt
    result = await agent_fn(task, candidate_id=0)
    score = await score_fn(result)
 
    if score >= threshold:
        return result  # Good enough on first try
 
    # Escalate: run 4 more candidates
    additional = await asyncio.gather(
        *[agent_fn(task, candidate_id=i) for i in range(1, 5)]
    )
 
    all_results = [(result, score)]
    for r in additional:
        s = await score_fn(r)
        all_results.append((r, s))
 
    all_results.sort(key=lambda x: x[1], reverse=True)
    return all_results[0][0]

class ProcessRewardScorer:
    """Score candidates at multiple checkpoints during execution."""
 
    def __init__(self, reward_model):
        self.reward_model = reward_model
        self.checkpoints = []
 
    async def score_checkpoint(self, agent_state: dict, step_name: str) -> float:
        """Score the agent's state at a specific checkpoint."""
        score = await self.reward_model.evaluate(
            step=step_name,
            state=agent_state,
        )
        self.checkpoints.append({
            "step": step_name,
            "score": score,
        })
        return score
 
    async def final_score(self) -> float:
        """Weighted combination of all checkpoint scores."""
        weights = {
            "problem_analysis": 0.15,
            "approach_selection": 0.20,
            "implementation": 0.25,
            "test_results": 0.40,
        }
        total = sum(
            cp["score"] * weights.get(cp["step"], 0.1)
            for cp in self.checkpoints
        )
        return total / sum(weights.values())

async def hierarchical_best_of_n(task: str):
    """Three-level hierarchical Best-of-N."""
 
    # Level 1: Strategy generation (N=5)
    strategies = await asyncio.gather(
        *[generate_strategy(task) for _ in range(5)]
    )
    strategy_scores = [(s, await score_strategy(s)) for s in strategies]
    strategy_scores.sort(key=lambda x: x[1], reverse=True)
    top_strategies = [s for s, _ in strategy_scores[:2]]  # Keep top 2
 
    # Level 2: Implementation (M=3 per strategy)
    implementations = []
    for strategy in top_strategies:
        impls = await asyncio.gather(
            *[implement_strategy(task, strategy) for _ in range(3)]
        )
        implementations.extend(impls)
 
    impl_scores = [(impl, await score_implementation(impl)) for impl in implementations]
    impl_scores.sort(key=lambda x: x[1], reverse=True)
    best_impl = impl_scores[0][0]
 
    # Level 3: Refinement (P=3)
    refined = await asyncio.gather(
        *[refine_implementation(best_impl, task) for _ in range(3)]
    )
    refined_scores = [(r, await score_implementation(r)) for r in refined]
    refined_scores.sort(key=lambda x: x[1], reverse=True)
 
    return refined_scores[0][0]

async def ensemble_selection(candidates: list, task: str) -> str:
    """Ask an LLM to synthesize the best parts of multiple candidates."""
    diffs = "\n\n---\n\n".join(
        f"Candidate {i+1}:\n{c['diff']}" for i, c in enumerate(candidates)
    )
 
    response = await call_llm(
        model="claude-sonnet-4",
        prompt=(
            f"Task: {task}\n\n"
            f"Here are {len(candidates)} candidate solutions:\n\n{diffs}\n\n"
            "Synthesize the best elements from these candidates into a single, "
            "optimal solution. Output only the unified diff."
        ),
    )
 
    return response

# Detect potential test overfitting
def check_for_overfitting(diff: str) -> float:
    """Return a penalty score for suspicious patterns."""
    suspicious_patterns = [
        r"if\s+.*==\s*['\"]specific_test_value['\"]",
        r"hardcoded",
        r"# TODO.*remove.*hack",
        r"return\s+\d+\s*$",  # Hardcoded return values
    ]
    import re
    penalty = sum(
        0.1 for pattern in suspicious_patterns
        if re.search(pattern, diff)
    )
    return min(penalty, 0.5)

def log_best_of_n_metrics(candidates: list, winner_idx: int):
    """Log metrics for Best-of-N analysis."""
    scores = [c["score"] for c in candidates]
    metrics = {
        "n": len(candidates),
        "winner_score": scores[winner_idx],
        "mean_score": sum(scores) / len(scores),
        "score_std": (sum((s - sum(scores)/len(scores))**2 for s in scores) / len(scores)) ** 0.5,
        "winning_margin": scores[winner_idx] - sorted(scores)[-2] if len(scores) > 1 else 0,
        "all_scores": scores,
        "num_passing": sum(1 for s in scores if s >= 0.9),
    }
    logger.info(f"Best-of-N metrics: {json.dumps(metrics)}")
    return metrics