"""
Complete pipeline for Best-of-N weighted selection on MATH-500.

This single script runs all steps:
1. Filter MATH-500 to 20 level 1-3 problems
2. Generate greedy (N=1) solutions and compute baseline accuracy
3. Sample N=16 solutions per problem with temperature sampling
4. Score all solutions with Skywork PRM (last-step prediction)
5. Compute weighted Best-of-N accuracy
6. Create dataset and push to HuggingFace Hub
7. Generate analysis plots and push them too

Reference papers:
- DeepMind (2408.03314): Scaling LLM Test-Time Compute, Section 5.1 + Appendix E
- Math-Shepherd (2312.08935): Process Reward Models, Section 3.4

Co-authored with Claude (Anthropic) as part of the HuggingFace internship exercise.
I can explain all code logic in detail.
"""

import json
import os
import random
import subprocess
import sys
import torch
import numpy as np
from collections import defaultdict
from typing import Optional

from datasets import Dataset, load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM


# ═══════════════════════════════════════════════════════════════════════════════
# Configuration
# ═══════════════════════════════════════════════════════════════════════════════
N_PROBLEMS = 20              # Number of problems to evaluate
N_SAMPLES = 16               # Number of solutions per problem for Best-of-N
TEMPERATURE = 0.7            # Sampling temperature for diverse solutions
MAX_NEW_TOKENS = 2048        # Max generation length
SEED = 42                    # Random seed for reproducibility
LLM_MODEL_ID = "Qwen/Qwen2.5-1.5B-Instruct"
PRM_MODEL_ID = "Skywork/Skywork-o1-Open-PRM-Qwen-2.5-1.5B"
DATASET_ID = "cmpatino/math500-bon-weighted-results"

OUTPUT_DIR = "/tmp/exercise_outputs"
os.makedirs(OUTPUT_DIR, exist_ok=True)

# System prompt: encourages chain-of-thought reasoning and \boxed{} format
SYSTEM_PROMPT = (
    "You are a helpful math assistant. Solve the problem step by step, "
    "showing your reasoning clearly. Put your final answer inside "
    "\\boxed{answer} at the end of your solution."
)


# ═══════════════════════════════════════════════════════════════════════════════
# Helper functions
# ═══════════════════════════════════════════════════════════════════════════════

def extract_boxed_solution(text: str) -> Optional[str]:
    """
    Extract content of the last \\boxed{} in text.
    Uses bracket-balanced parsing for nested braces.
    Source: https://gist.github.com/lewtun/9c2ce1937b741404090a3dc4c7c022b3
    """
    try:
        start_index = text.rindex("\\boxed{")
        content_start = start_index + 7
        bracket_count = 1
        current_pos = content_start
        while bracket_count > 0 and current_pos < len(text):
            if text[current_pos] == "{":
                bracket_count += 1
            elif text[current_pos] == "}":
                bracket_count -= 1
            current_pos += 1
        if bracket_count == 0:
            return text[content_start : current_pos - 1].strip()
        return None
    except (ValueError, Exception):
        return None


def weighted_best_of_n(extracted_answers, prm_scores):
    """
    Weighted Best-of-N selection (DeepMind 2408.03314, Eq. from Section 5.1):
    â = argmax_a  Σᵢ 𝟙(aᵢ = a) · score(sᵢ)

    Groups solutions by final answer, sums their PRM scores,
    and selects the answer group with the highest total.
    """
    answer_scores = defaultdict(float)
    for answer, score in zip(extracted_answers, prm_scores):
        if answer is None:
            continue  # Skip unparseable solutions
        answer_scores[answer] += score
    if not answer_scores:
        return None, {}
    best_answer = max(answer_scores, key=answer_scores.get)
    return best_answer, dict(answer_scores)


def standard_best_of_n(extracted_answers, prm_scores):
    """Standard Best-of-N: pick the single solution with highest PRM score."""
    best_idx, best_score = None, -1
    for i, (answer, score) in enumerate(zip(extracted_answers, prm_scores)):
        if answer is not None and score > best_score:
            best_score = score
            best_idx = i
    return extracted_answers[best_idx] if best_idx is not None else None


def majority_vote(extracted_answers):
    """Pure majority vote: pick the most frequent answer."""
    counts = defaultdict(int)
    for answer in extracted_answers:
        if answer is not None:
            counts[answer] += 1
    return max(counts, key=counts.get) if counts else None


# ═══════════════════════════════════════════════════════════════════════════════
# STEP 1: Filter MATH-500 to level 1-3 problems
# ═══════════════════════════════════════════════════════════════════════════════
print("=" * 70)
print("STEP 1: Loading and filtering MATH-500 dataset")
print("=" * 70)

dataset = load_dataset("HuggingFaceH4/MATH-500", split="test")
print(f"Total problems: {len(dataset)}")

# Filter to levels 1-3 — these are easier problems that a 1.5B model
# can reasonably attempt. Levels 4-5 are too hard for small models.
filtered = dataset.filter(lambda x: x["level"] in [1, 2, 3])
print(f"Level 1-3 problems: {len(filtered)}")

# Fixed random sample for reproducibility
random.seed(SEED)
indices = random.sample(range(len(filtered)), k=N_PROBLEMS)
problems = filtered.select(indices)

problems_data = []
for i, p in enumerate(problems):
    problems_data.append({
        "idx": i,
        "problem": p["problem"],
        "solution": p["solution"],
        "answer": p["answer"],
        "subject": p["subject"],
        "level": p["level"],
        "unique_id": p["unique_id"],
    })
    print(f"  [{i+1:2d}] L{p['level']} {p['subject']:25s} {p['unique_id']}")

# Save for reference
with open(os.path.join(OUTPUT_DIR, "filtered_problems.json"), "w") as f:
    json.dump(problems_data, f, indent=2)


# ═══════════════════════════════════════════════════════════════════════════════
# STEP 2: Generate greedy (N=1) solutions
# ═══════════════════════════════════════════════════════════════════════════════
print("\n" + "=" * 70)
print("STEP 2: Generating greedy solutions (N=1)")
print("=" * 70)

tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    LLM_MODEL_ID, torch_dtype=torch.bfloat16, device_map="auto"
)


def generate_batch(problems_data, model, tokenizer, n, do_sample, temperature=None):
    """Generate n solutions per problem. Returns list of solution lists."""
    all_solutions = []
    for i, p in enumerate(problems_data):
        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": p["problem"]},
        ]
        prompt = tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

        solutions = []
        for j in range(n):
            gen_kwargs = {"max_new_tokens": MAX_NEW_TOKENS, "do_sample": do_sample}
            if do_sample and temperature:
                gen_kwargs["temperature"] = temperature
                gen_kwargs["top_p"] = 0.95
            with torch.no_grad():
                output = model.generate(**inputs, **gen_kwargs)
            generated = output[0][inputs["input_ids"].shape[1]:]
            solutions.append(tokenizer.decode(generated, skip_special_tokens=True))

        all_solutions.append(solutions)
        ans = extract_boxed_solution(solutions[0]) if n == 1 else "..."
        tag = "greedy" if n == 1 else f"N={n}"
        print(f"  [{i+1:2d}/{len(problems_data)}] {tag} | {p['unique_id']} | answer={ans}")

    return all_solutions


# Greedy decoding (N=1, deterministic)
greedy_solutions = generate_batch(problems_data, model, tokenizer, n=1, do_sample=False)

# Evaluate greedy accuracy
greedy_correct = 0
for p, sols in zip(problems_data, greedy_solutions):
    extracted = extract_boxed_solution(sols[0])
    p["greedy_solution"] = sols[0]
    p["greedy_extracted_answer"] = extracted
    p["greedy_correct"] = (extracted is not None) and (extracted == p["answer"])
    if p["greedy_correct"]:
        greedy_correct += 1
    status = "✓" if p["greedy_correct"] else "✗"
    print(f"  {status} Expected: {p['answer']:20s} | Got: {str(extracted):20s} | {p['unique_id']}")

greedy_acc = greedy_correct / len(problems_data)
print(f"\n>>> Greedy accuracy: {greedy_correct}/{len(problems_data)} = {greedy_acc:.0%}")


# ═══════════════════════════════════════════════════════════════════════════════
# STEP 3: Sample N=16 solutions per problem
# ═══════════════════════════════════════════════════════════════════════════════
print("\n" + "=" * 70)
print(f"STEP 3: Sampling N={N_SAMPLES} solutions per problem (T={TEMPERATURE})")
print("=" * 70)

sampled_solutions = generate_batch(
    problems_data, model, tokenizer,
    n=N_SAMPLES, do_sample=True, temperature=TEMPERATURE
)

# Save solutions and free LLM memory
for p, sols in zip(problems_data, sampled_solutions):
    p["sampled_solutions"] = sols

with open(os.path.join(OUTPUT_DIR, "sampled_solutions.json"), "w") as f:
    json.dump(problems_data, f, indent=2)

del model
torch.cuda.empty_cache()
print("Freed LLM memory for PRM loading.")


# ═══════════════════════════════════════════════════════════════════════════════
# STEP 4: Score with Skywork PRM (last-step prediction)
# ═══════════════════════════════════════════════════════════════════════════════
print("\n" + "=" * 70)
print("STEP 4: Scoring solutions with Skywork PRM")
print("=" * 70)

# Clone the Skywork PRM inference repo for the custom PRM_MODEL class
PRM_REPO_PATH = "/tmp/skywork-o1-prm-inference"
if not os.path.exists(PRM_REPO_PATH):
    print("Cloning Skywork PRM inference repo...")
    subprocess.run(
        ["git", "clone", "https://github.com/SkyworkAI/skywork-o1-prm-inference.git", PRM_REPO_PATH],
        check=True,
    )
sys.path.insert(0, PRM_REPO_PATH)

from model_utils.prm_model import PRM_MODEL
from model_utils.io_utils import prepare_input, prepare_batch_input_for_model, derive_step_rewards

prm_tokenizer = AutoTokenizer.from_pretrained(PRM_MODEL_ID, trust_remote_code=True)
prm_model = PRM_MODEL.from_pretrained(PRM_MODEL_ID, device_map="auto").eval()
prm_device = next(prm_model.pretrained_model.parameters()).device
print(f"PRM loaded on {prm_device}")


def score_solution(problem: str, solution: str) -> float:
    """
    Score a single solution using the PRM's last-step prediction.

    Per DeepMind (2408.03314, Appendix E): "We use the PRM's prediction at the
    last step as the full-answer score" — this outperforms min/product aggregation
    when the PRM is trained with soft MC-return labels.

    Returns: float in [0, 1] — the sigmoid-normalized score at the last step.
    """
    input_ids, steps, reward_flags = prepare_input(
        problem, solution, prm_tokenizer, step_token="\n"
    )
    input_ids_t, attn_mask_t, flags_t = prepare_batch_input_for_model(
        [input_ids], [reward_flags], prm_tokenizer.pad_token_id
    )
    input_ids_t = input_ids_t.to(prm_device)
    attn_mask_t = attn_mask_t.to(prm_device)
    flags_t = flags_t.to(prm_device)

    with torch.no_grad():
        _, _, rewards = prm_model(
            input_ids=input_ids_t, attention_mask=attn_mask_t, return_probs=True
        )
    step_rewards = derive_step_rewards(rewards, flags_t)
    # Return last step score (or 0.0 if no steps found)
    return step_rewards[0][-1] if step_rewards[0] else 0.0


# Score all sampled solutions
for i, p in enumerate(problems_data):
    print(f"\n  Scoring problem {i+1}/{len(problems_data)}: {p['unique_id']}")
    scores = []
    extracted_answers = []
    for j, sol in enumerate(p["sampled_solutions"]):
        score = score_solution(p["problem"], sol)
        scores.append(score)
        extracted_answers.append(extract_boxed_solution(sol))
        if (j + 1) % 8 == 0:
            print(f"    Scored {j+1}/{N_SAMPLES} (last: {score:.4f})")
    p["prm_scores"] = scores
    p["extracted_answers"] = extracted_answers

# Save scored results
with open(os.path.join(OUTPUT_DIR, "scored_results.json"), "w") as f:
    json.dump(problems_data, f, indent=2)

del prm_model
torch.cuda.empty_cache()


# ═══════════════════════════════════════════════════════════════════════════════
# STEP 5: Compute Best-of-N with weighted selection
# ═══════════════════════════════════════════════════════════════════════════════
print("\n" + "=" * 70)
print("STEP 5: Computing Best-of-N accuracy")
print("=" * 70)

weighted_correct = 0
standard_correct = 0
majority_correct_count = 0

bon_summary = []
for p in problems_data:
    gt = p["answer"]

    # Weighted BoN
    w_ans, w_scores = weighted_best_of_n(p["extracted_answers"], p["prm_scores"])
    w_ok = (w_ans is not None) and (w_ans == gt)
    if w_ok: weighted_correct += 1

    # Standard BoN
    s_ans = standard_best_of_n(p["extracted_answers"], p["prm_scores"])
    s_ok = (s_ans is not None) and (s_ans == gt)
    if s_ok: standard_correct += 1

    # Majority vote
    m_ans = majority_vote(p["extracted_answers"])
    m_ok = (m_ans is not None) and (m_ans == gt)
    if m_ok: majority_correct_count += 1

    n_correct = sum(1 for a in p["extracted_answers"] if a == gt)

    bon_summary.append({
        "unique_id": p["unique_id"],
        "level": p["level"],
        "subject": p["subject"],
        "ground_truth": gt,
        "greedy_answer": p["greedy_extracted_answer"],
        "greedy_correct": p["greedy_correct"],
        "weighted_bon_answer": w_ans,
        "weighted_bon_correct": w_ok,
        "standard_bon_answer": s_ans,
        "standard_bon_correct": s_ok,
        "majority_vote_answer": m_ans,
        "majority_vote_correct": m_ok,
        "n_correct_in_16": n_correct,
        "answer_score_breakdown": w_scores,
        "prm_scores": p["prm_scores"],
    })

    sg = "✓" if p["greedy_correct"] else "✗"
    sw = "✓" if w_ok else "✗"
    print(f"  {sg}→{sw} | {p['unique_id']:40s} | GT={gt:15s} | Greedy={str(p['greedy_extracted_answer']):15s} | WBoN={str(w_ans):15s} | {n_correct}/16 correct")

n = len(problems_data)
greedy_total = sum(1 for p in problems_data if p["greedy_correct"])
print(f"\n{'='*70}")
print(f"RESULTS SUMMARY")
print(f"{'='*70}")
print(f"  Greedy (N=1):              {greedy_total}/{n} = {greedy_total/n:.0%}")
print(f"  Majority Vote (N=16):      {majority_correct_count}/{n} = {majority_correct_count/n:.0%}")
print(f"  Standard Best-of-N (N=16): {standard_correct}/{n} = {standard_correct/n:.0%}")
print(f"  Weighted Best-of-N (N=16): {weighted_correct}/{n} = {weighted_correct/n:.0%}")

with open(os.path.join(OUTPUT_DIR, "bon_results.json"), "w") as f:
    json.dump(bon_summary, f, indent=2)


# ═══════════════════════════════════════════════════════════════════════════════
# STEP 5b: Accuracy vs N analysis
# ═══════════════════════════════════════════════════════════════════════════════
print("\n" + "=" * 70)
print("ANALYSIS: Accuracy vs N")
print("=" * 70)

random.seed(SEED)
n_values = [1, 2, 4, 8, 16]
n_trials = 50

accuracy_by_n = {}
for n_val in n_values:
    if n_val == 16:
        correct = sum(1 for p in problems_data
                      for _ in [weighted_best_of_n(p["extracted_answers"], p["prm_scores"])]
                      if _[0] == p["answer"])
        acc = correct / len(problems_data)
    else:
        trial_accs = []
        for _ in range(n_trials):
            correct = 0
            for p in problems_data:
                idx = random.sample(range(16), n_val)
                sub_a = [p["extracted_answers"][j] for j in idx]
                sub_s = [p["prm_scores"][j] for j in idx]
                ans, _ = weighted_best_of_n(sub_a, sub_s)
                if ans == p["answer"]:
                    correct += 1
            trial_accs.append(correct / len(problems_data))
        acc = sum(trial_accs) / len(trial_accs)
    accuracy_by_n[n_val] = acc
    print(f"  N={n_val:2d}: {acc:.1%}")

with open(os.path.join(OUTPUT_DIR, "accuracy_by_n.json"), "w") as f:
    json.dump(accuracy_by_n, f, indent=2)


# ═══════════════════════════════════════════════════════════════════════════════
# STEP 6: Generate plots
# ═══════════════════════════════════════════════════════════════════════════════
print("\n" + "=" * 70)
print("STEP 6: Generating analysis plots")
print("=" * 70)

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import Patch

plt.rcParams.update({"font.size": 11, "figure.dpi": 150})

# --- Plot 1: Overall accuracy comparison ---
fig, ax = plt.subplots(figsize=(8, 5))
methods = ["Greedy\n(N=1)", "Majority Vote\n(N=16)", "Standard BoN\n(N=16)", "Weighted BoN\n(N=16)"]
accs = [
    greedy_total / n,
    majority_correct_count / n,
    standard_correct / n,
    weighted_correct / n,
]
colors = ["#4C72B0", "#55A868", "#C44E52", "#8172B2"]
bars = ax.bar(methods, accs, color=colors, edgecolor="white", linewidth=1.5)
for bar, a in zip(bars, accs):
    ax.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.01,
            f"{a:.0%}", ha="center", va="bottom", fontweight="bold", fontsize=12)
ax.set_ylabel("Accuracy")
ax.set_title("Math Problem Accuracy: Greedy vs Best-of-N Methods\n(20 MATH-500 problems, Levels 1-3)")
ax.set_ylim(0, 1.15)
ax.grid(axis="y", alpha=0.3)
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "plot1_accuracy_comparison.png"))
plt.close()

# --- Plot 2: Accuracy vs N ---
fig, ax = plt.subplots(figsize=(7, 5))
ns = sorted(accuracy_by_n.keys())
acc_vals = [accuracy_by_n[nv] for nv in ns]
ax.plot(ns, acc_vals, "o-", color="#8172B2", linewidth=2, markersize=8, label="Weighted BoN")
ax.axhline(y=greedy_total/n, color="#4C72B0", linestyle="--", linewidth=1.5,
           label=f"Greedy baseline ({greedy_total/n:.0%})")
for nv, a in zip(ns, acc_vals):
    ax.annotate(f"{a:.0%}", (nv, a), textcoords="offset points", xytext=(0, 10), ha="center")
ax.set_xlabel("N (number of samples)")
ax.set_ylabel("Accuracy")
ax.set_title("Weighted Best-of-N Accuracy vs Number of Samples")
ax.set_xticks(ns)
ax.set_ylim(0, 1.1)
ax.legend()
ax.grid(alpha=0.3)
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "plot2_accuracy_vs_n.png"))
plt.close()

# --- Plot 3: Per-problem analysis ---
fig, ax = plt.subplots(figsize=(12, 5))
cat_colors = {
    "Both correct": "#55A868", "Only BoN correct": "#8172B2",
    "Only Greedy correct": "#C44E52", "Both wrong": "#CCCCCC"
}
bar_colors = []
for s in bon_summary:
    g, b = s["greedy_correct"], s["weighted_bon_correct"]
    if g and b: bar_colors.append(cat_colors["Both correct"])
    elif not g and b: bar_colors.append(cat_colors["Only BoN correct"])
    elif g and not b: bar_colors.append(cat_colors["Only Greedy correct"])
    else: bar_colors.append(cat_colors["Both wrong"])

x = range(len(bon_summary))
heights = [s["n_correct_in_16"] for s in bon_summary]
ax.bar(x, heights, color=bar_colors, edgecolor="white", linewidth=0.5)
ax.set_xticks(x)
labels = [f"L{s['level']}: {s['unique_id'].split('/')[-1].replace('.json','')[:12]}" for s in bon_summary]
ax.set_xticklabels(labels, rotation=45, ha="right", fontsize=8)
ax.set_ylabel("# Correct Solutions (out of 16)")
ax.set_title("Per-Problem: Correct Solutions in N=16 Sample")
legend_elements = [Patch(facecolor=c, label=l) for l, c in cat_colors.items()]
ax.legend(handles=legend_elements, loc="upper right", fontsize=9)
ax.grid(axis="y", alpha=0.3)
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "plot3_per_problem.png"))
plt.close()

# --- Plot 4: PRM score distribution ---
fig, ax = plt.subplots(figsize=(7, 5))
correct_scores, incorrect_scores = [], []
for p in problems_data:
    for ans, sc in zip(p["extracted_answers"], p["prm_scores"]):
        (correct_scores if ans == p["answer"] else incorrect_scores).append(sc)

bins = np.linspace(0, 1, 25)
ax.hist(correct_scores, bins=bins, alpha=0.7, label=f"Correct ({len(correct_scores)})", color="#55A868")
ax.hist(incorrect_scores, bins=bins, alpha=0.7, label=f"Incorrect ({len(incorrect_scores)})", color="#C44E52")
ax.set_xlabel("PRM Last-Step Score")
ax.set_ylabel("Count")
ax.set_title("PRM Score Distribution: Correct vs Incorrect Solutions")
ax.legend()
ax.grid(alpha=0.3)
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "plot4_prm_scores.png"))
plt.close()

print("All plots saved.")


# ═══════════════════════════════════════════════════════════════════════════════
# STEP 7: Push dataset to Hub
# ═══════════════════════════════════════════════════════════════════════════════
print("\n" + "=" * 70)
print("STEP 7: Pushing dataset to HuggingFace Hub")
print("=" * 70)

rows = []
for p, s in zip(problems_data, bon_summary):
    rows.append({
        "problem": p["problem"],
        "ground_truth_solution": p["solution"],
        "ground_truth_answer": p["answer"],
        "subject": p["subject"],
        "level": p["level"],
        "unique_id": p["unique_id"],
        "greedy_solution": p["greedy_solution"],
        "greedy_extracted_answer": p["greedy_extracted_answer"],
        "greedy_correct": p["greedy_correct"],
        "bon_weighted_answer": s["weighted_bon_answer"],
        "bon_weighted_correct": s["weighted_bon_correct"],
        "bon_standard_answer": s["standard_bon_answer"],
        "bon_standard_correct": s["standard_bon_correct"],
        "bon_majority_answer": s["majority_vote_answer"],
        "bon_majority_correct": s["majority_vote_correct"],
        "sampled_solutions": p["sampled_solutions"],
        "sampled_extracted_answers": p["extracted_answers"],
        "sampled_prm_scores": p["prm_scores"],
        "n_correct_in_16": s["n_correct_in_16"],
        "answer_score_breakdown": json.dumps(s["answer_score_breakdown"]),
    })

hf_dataset = Dataset.from_list(rows)
hf_dataset.push_to_hub(DATASET_ID, split="test")
print(f"Dataset pushed to: https://huggingface.co/datasets/{DATASET_ID}")

# Also upload the plots as artifacts
from huggingface_hub import HfApi
api = HfApi()
for plot_file in ["plot1_accuracy_comparison.png", "plot2_accuracy_vs_n.png",
                   "plot3_per_problem.png", "plot4_prm_scores.png"]:
    plot_path = os.path.join(OUTPUT_DIR, plot_file)
    if os.path.exists(plot_path):
        api.upload_file(
            path_or_fileobj=plot_path,
            path_in_repo=f"plots/{plot_file}",
            repo_id=DATASET_ID,
            repo_type="dataset",
        )
        print(f"  Uploaded {plot_file}")

# Upload the results JSON files too
for json_file in ["filtered_problems.json", "bon_results.json", "accuracy_by_n.json"]:
    json_path = os.path.join(OUTPUT_DIR, json_file)
    if os.path.exists(json_path):
        api.upload_file(
            path_or_fileobj=json_path,
            path_in_repo=f"results/{json_file}",
            repo_id=DATASET_ID,
            repo_type="dataset",
        )
        print(f"  Uploaded {json_file}")


# ═══════════════════════════════════════════════════════════════════════════════
# Final summary
# ═══════════════════════════════════════════════════════════════════════════════
print("\n" + "=" * 70)
print("FINAL RESULTS")
print("=" * 70)
print(f"  Greedy (N=1):              {greedy_total}/{len(problems_data)} = {greedy_total/len(problems_data):.0%}")
print(f"  Majority Vote (N=16):      {majority_correct_count}/{len(problems_data)} = {majority_correct_count/len(problems_data):.0%}")
print(f"  Standard Best-of-N (N=16): {standard_correct}/{len(problems_data)} = {standard_correct/len(problems_data):.0%}")
print(f"  Weighted Best-of-N (N=16): {weighted_correct}/{len(problems_data)} = {weighted_correct/len(problems_data):.0%}")
print(f"\n  Dataset: https://huggingface.co/datasets/{DATASET_ID}")
print("=" * 70)
print("DONE!")