# /// script
# dependencies = ["datasets", "huggingface_hub>=0.34"]
# ///
"""Generate SDPO hints for the MATH dataset (hard levels) with a large model on an HF Inference Endpoint.

Pushes two datasets built from the SAME problems:
  - <output_repo>       : problem, answer, solution, privileged_context   (for A hints / B baseline)
  - <output_repo>-plain : problem, answer, solution                       (for C gold_feedback)

The large model only *generates* text (no logprobs), so any standard IE chat deployment works.

Usage:
    HF_TOKEN=hf_... python generate_hints_math.py \
        --endpoint_url https://xxxx.endpoints.huggingface.cloud \
        --output_repo sergiopaniego/math-sdpo-hints \
        --num_examples 600 --min_level 4
"""
import argparse
import os
import time
from concurrent.futures import ThreadPoolExecutor

from datasets import load_dataset
from huggingface_hub import InferenceClient


HINT_SYSTEM = (
    "You are a math expert. Given a competition math problem, write a concise worked solution "
    "(the key steps, not verbose) and end with the final answer as \\boxed{<answer>}."
)


def make_hint(client, problem, model, max_retries=3):
    for attempt in range(max_retries):
        try:
            resp = client.chat_completion(
                messages=[
                    {"role": "system", "content": HINT_SYSTEM},
                    {"role": "user", "content": problem},
                ],
                model=model,
                max_tokens=512,
                temperature=0.3,
            )
            return resp.choices[0].message.content.strip()
        except Exception:
            if attempt == max_retries - 1:
                return ""
            time.sleep(2 * (attempt + 1))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--endpoint_url", required=True, help="IE endpoint base URL")
    ap.add_argument("--model", default=None, help="Model id for chat_completion (optional on a dedicated endpoint)")
    ap.add_argument("--num_examples", type=int, default=600)
    ap.add_argument("--min_level", type=int, default=4, help="Keep only problems with level >= this (1-5).")
    ap.add_argument("--output_repo", required=True, help="Hub dataset repo id (hints version).")
    ap.add_argument("--concurrency", type=int, default=8)
    args = ap.parse_args()

    client = InferenceClient(base_url=args.endpoint_url, token=os.environ.get("HF_TOKEN"))

    ds = load_dataset("nlile/hendrycks-MATH-benchmark", split="train")
    ds = ds.filter(lambda ex: ex["level"] >= args.min_level)
    ds = ds.shuffle(seed=42).select(range(min(args.num_examples, len(ds))))
    print(f"Selected {len(ds)} problems (level >= {args.min_level}).")

    t0 = time.time()
    with ThreadPoolExecutor(max_workers=args.concurrency) as pool:
        hints = list(pool.map(lambda ex: make_hint(client, ex["problem"], args.model), ds))

    n_empty = sum(1 for h in hints if not h)
    ds = ds.add_column("privileged_context", hints).filter(lambda ex: bool(ex["privileged_context"]))

    ds.push_to_hub(args.output_repo)
    ds.remove_columns("privileged_context").push_to_hub(f"{args.output_repo}-plain")
    print(f"Pushed {len(ds)} examples to {args.output_repo} (+ -plain) ({n_empty} empty dropped) in {time.time() - t0:.0f}s")
    print("Columns:", ds.column_names)


if __name__ == "__main__":
    main()