"""02 APE, Automatic Prompt Engineer. The model reverse-engineers the instruction from examples,
then a second pass scores each candidate and keeps the best.

Paper: Zhou et al. 2022, arxiv.org/abs/2211.01910 (human-level instructions on 24/24 tasks).
The generation prompt below is the paper's actual template ("I gave a friend an instruction...").
ES: El modelo deduce la instruccion a partir de ejemplos, luego un juez califica cada candidata.

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

MODEL = os.environ.get("METAPROMPT_MODEL", "haiku")
N_CANDIDATES = 4

# Swap these for YOUR task. Train pairs teach the instruction; test pairs score it.
TRAIN = [
    ("The meeting is at 3pm.", "La reu es a las 3pm."),
    ("Can you send me the file?", "Me mandas el archivo?"),
    ("I will be 10 minutes late.", "Llego 10 min tarde."),
    ("That idea is great.", "Esa idea esta buenisima."),
]
TEST = [
    ("See you tomorrow at the office.", "Nos vemos manana en la oficina."),
    ("I did not understand the email.", "No entendi el correo."),
]


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 main():
    pairs = "\n".join(f"Input: {i}\nOutput: {o}" for i, o in TRAIN)

    # Step 1 - generate candidates with APE's own template.
    gen = llm(
        "I gave a friend an instruction and four inputs. The friend read the instruction and "
        "wrote an output for every one of the inputs. Here are the input-output pairs:\n\n"
        f"{pairs}\n\n"
        f"Write {N_CANDIDATES} different guesses for what the instruction was. "
        'Reply ONLY with a JSON array of strings, like ["instruction one", "instruction two"].'
    )
    m = re.search(r"\[.*\]", gen, re.S)
    candidates = json.loads(m.group(0)) if m else [gen]
    print("== CANDIDATE INSTRUCTIONS ==")
    for i, c in enumerate(candidates):
        print(f"  [{i}] {c}")

    # Step 2 - score each candidate on held-out pairs.
    print("\n== SCORING ON HELD-OUT PAIRS ==")
    scored = []
    for i, cand in enumerate(candidates):
        outs = []
        for inp, _ in TEST:
            outs.append(llm(f"Instruction: {cand}\nInput: {inp}\nOutput only the result."))
        judge = llm(
            "Score how well the actual outputs match the expected outputs, 0 to 10 overall. "
            "Reply with ONLY the number.\n\n"
            + "\n".join(
                f"Input: {t[0]}\nExpected: {t[1]}\nActual: {o}" for t, o in zip(TEST, outs)
            )
        )
        score = float(re.search(r"\d+(\.\d+)?", judge).group(0)) if re.search(r"\d", judge) else 0.0
        scored.append((score, cand))
        print(f"  [{i}] score {score:>4}  {cand[:70]}")

    best = max(scored)
    print(f"\n== WINNER ==\n{best[1]}\n(score {best[0]}/10)")
    print("\nThe best prompt in your repo was probably found, not written. / Encontrada, no escrita.")


if __name__ == "__main__":
    main()