"""03 OPRO, Optimization by PROmpting. The whole score history becomes the next prompt.
The model reads every past instruction with its measured score, sorted low to high, and
proposes one it expects to score higher. This is how "Take a deep breath" was found (80.2% GSM8K).

Paper: Yang et al. 2023, arxiv.org/abs/2309.03409. <INS> convention is the paper's own.
ES: El historial de puntajes ES el siguiente prompt. Asi se encontro "Take a deep breath".

Note: on a strong model this tiny eval set may already score perfect at round 0. Watch the
trajectory anyway: clever-sounding candidates often score WORSE, and catching a regression
is the same muscle as finding a gain. The optimizer is only as good as the scorer you give it.
ES: A veces el candidato "inteligente" puntua PEOR. Detectar retrocesos es el mismo musculo.

Run:  python 03_opro.py
Needs: Claude Code installed (the `claude` CLI on your PATH). No API keys.
"""
import os, re, shutil, subprocess, sys

MODEL = os.environ.get("METAPROMPT_MODEL", "haiku")
ROUNDS = 3

# Tiny eval set with known answers. Trick problems on purpose: a bare instruction
# misses some, which gives the optimizer room to climb. Swap for YOUR task + metric.
PROBLEMS = [
    ("A snail climbs 3 meters each day and slips back 2 meters each night. The well is 10 meters deep. On which day number does it reach the top?", "8"),
    ("A clock takes 5 seconds to chime 6 times, with evenly spaced chimes. How many seconds to chime 12 times?", "11"),
    ("A farmer has 17 sheep. All but 9 run away. How many sheep are left?", "9"),
    ("Ana has twice as many stickers as Luis. Together they have 36. How many MORE stickers does Ana have than Luis?", "12"),
    ("A bus had 23 people. 9 got off, 4 got on. How many people are on the bus?", "18"),
]


def llm(prompt: str) -> str:
    exe = shutil.which("claude")
    if not exe:
        sys.exit("claude CLI not found. Install Claude Code first: https://claude.com/claude-code")
    r = subprocess.run([exe, "-p", prompt, "--model", MODEL], capture_output=True, text=True, encoding="utf-8")
    return (r.stdout or "").strip()


def score(instruction: str) -> int:
    """Run the instruction on every problem; count exact final answers."""
    correct = 0
    for q, ans in PROBLEMS:
        out = llm(f"Q: {q}\nA: {instruction}\nEnd your reply with: FINAL: <number>")
        m = re.search(r"FINAL:\s*(-?\d+)", out)
        if m and m.group(1) == ans:
            correct += 1
    return correct


def main():
    history = [("Solve the problem.", score("Solve the problem."))]
    print(f"round 0 · '{history[0][0]}' -> {history[0][1]}/{len(PROBLEMS)}")

    for rnd in range(1, ROUNDS + 1):
        # OPRO's core: the meta-prompt carries the full scored trajectory, worst first.
        trajectory = "\n".join(
            f"text: {ins}\nscore: {sc}" for ins, sc in sorted(history, key=lambda x: x[1])
        )
        proposal = llm(
            "Your task is to generate the instruction <INS>. Below are previous instructions "
            "with their scores (number of problems solved out of "
            f"{len(PROBLEMS)}). Higher is better.\n\n{trajectory}\n\n"
            "The instruction is prepended to the answer of a math word problem. "
            "Generate ONE new instruction, different from all above, that you expect to score "
            "higher. Reply with only the instruction text."
        ).strip().strip('"')
        sc = score(proposal)
        history.append((proposal, sc))
        print(f"round {rnd} · '{proposal[:70]}' -> {sc}/{len(PROBLEMS)}")

    best = max(history, key=lambda x: x[1])
    print(f"\n== BEST INSTRUCTION ==\n{best[0]}  ({best[1]}/{len(PROBLEMS)})")
    print("\nOptimization needs a scorer more than it needs a genius. / Necesitas un medidor, no un genio.")


if __name__ == "__main__":
    main()