#!/usr/bin/env python
"""End-to-end pass@1 of the released model + search harness on miniF2F.

This is the script behind the documented 30/89 = 33.7% (valid, first 89 problems).
It needs a working **Lean 4 + mathlib4** installation (see docs/USAGE.md), because
every candidate is verified by the Lean kernel - that is the ground truth here.

  python eval_minif2f.py --n 89 --ckpt checkpoints/stage3-e3            # reproduce 30/89
  python eval_minif2f.py --n 244 --ckpt checkpoints/e6 --lean-parallel 6
Results are appended to reports/search_eval_<split>[_tag].jsonl so runs can be extended;
--start lets you resume an interrupted run.

NOTE this harness is what makes the pass@1 numbers non-trivial: the model alone proposes
candidates, the injection table below plus depth-batched search do the closing. Both are
shipped here verbatim so the number can be reproduced without the training code.
"""
import argparse
import json
import os
import sys
import time

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.abspath(os.path.join(HERE, os.pardir))
sys.path.insert(0, HERE)
from leanoar.search import ProofSearch                  # noqa: E402

FORCE_SETS = {
    # v1: the original 14 entries (kept for historical comparison; ring_nf / positivity /
    #     field_simp / decide / bare norm_num never closed anything in the 92 solved proofs)
    'v1': ['rfl', 'simp', 'simp_all', 'ring', 'ring_nf', 'norm_num', 'omega', 'linarith',
           'nlinarith', 'positivity', 'field_simp', 'trivial', 'aesop', 'decide'],
    # v2 (used by the documented results): 5 dead entries dropped, measured combinations added
    #     (`simp_all ; ring` closed 6 goals, `simp_all ; norm_num` 2), plus gcongr for
    #     inequality goals. 16 entries ~= v1's 14, so cost is unchanged and the table content
    #     is the single variable.
    'v2': ['rfl', 'ring', 'omega', 'simp_all', 'linarith', 'trivial', 'nlinarith', 'aesop',
           'simp', 'norm_num',
           'simp_all ; ring', 'simp_all ; norm_num', 'simp_all ; omega', 'simp_all ; linarith',
           'gcongr', 'field_simp ; ring'],
}

ap = argparse.ArgumentParser()
ap.add_argument('--split', default='valid')
ap.add_argument('--n', type=int, default=89)
ap.add_argument('--start', type=int, default=0)
ap.add_argument('--k', type=int, default=16)
ap.add_argument('--beam', type=int, default=4)
ap.add_argument('--depth', type=int, default=6)
ap.add_argument('--budget', type=float, default=45.0)
ap.add_argument('--round-budget', type=float, default=30.0)
ap.add_argument('--verify-chunk', type=int, default=24)
ap.add_argument('--force', action='store_true', help='inject the automation table')
ap.add_argument('--force-depth', type=int, default=99, help='max depth at which to inject it')
ap.add_argument('--force-set', default='v2', choices=['v1', 'v2'])
ap.add_argument('--lean-parallel', type=int, default=6, help='concurrent lean --json chunks')
ap.add_argument('--narrow-imports', action='store_true',
                help='replace the full Mathlib import per problem by #min_imports + Mathlib.Tactic '
                     '(measured 3.41s -> 1.74s per lean call); failures that look like missing '
                     'names are automatically re-checked with the full import')
ap.add_argument('--no-model', action='store_true',
                help='ablation: verify only the injected automation table, no model proposals '
                     '(measures what the harness alone achieves)')
ap.add_argument('--ckpt', default='checkpoints/stage3-e3')
ap.add_argument('--tag', default='')
ap.add_argument('--outdir', default=os.path.join(os.getcwd(), 'reports'))
A = ap.parse_args()

rows = [json.loads(l) for l in open(os.path.join(ROOT, 'data/minif2f', f'{A.split}.jsonl'))]
todo = rows[A.start:A.start + A.n]
suffix = f'_{A.tag}' if A.tag else ''
out_path = os.path.join(A.outdir, f'search_eval_{A.split}{suffix}.jsonl')
os.makedirs(A.outdir, exist_ok=True)

ckpt = A.ckpt if os.path.isabs(A.ckpt) else os.path.join(ROOT, A.ckpt)
print('loading model + verifier ...', flush=True)
t0 = time.time()
FORCE = FORCE_SETS[A.force_set] if A.force else None
S = ProofSearch(ckpt=ckpt, k=A.k, beam=A.beam, max_depth=A.depth, time_budget=A.budget,
                round_budget=A.round_budget, verify_chunk=A.verify_chunk,
                force_tactics=FORCE, force_depth=A.force_depth,
                lean_parallel=A.lean_parallel,
                narrow_imports=A.narrow_imports, model_proposals=not A.no_model,
                tok_file=os.path.join(ROOT, 'tokenizer_v1.json'),
                steps_root=os.path.join(ROOT, 'data/whitelist'))
print(f'ready in {time.time() - t0:.1f}s | {len(todo)} problems '
      f'| k={A.k} beam={A.beam} depth={A.depth} budget={A.budget}s '
      f'| force={A.force_set if A.force else "off"} depth<={A.force_depth}', flush=True)

solved = 0
t_start = time.time()
with open(out_path, 'a') as fout:
    for i, r in enumerate(todo):
        t1 = time.time()
        res = S.solve(r['header'], r['statement'], max_depth=A.depth, time_budget=A.budget)
        rec = {'id': r['id'], 'solved': res['solved'], 'depth': res.get('depth'),
               'seconds': round(time.time() - t1, 1), 'verified': res.get('verified', 0),
               'proof': res.get('proof'), 'confirmed': (res.get('confirm') or {}).get('ok'),
               'reason': res.get('reason')}
        fout.write(json.dumps(rec, ensure_ascii=False) + '\n')
        fout.flush()
        solved += res['solved']
        done = i + 1
        el = time.time() - t_start
        print(f'[{done}/{len(todo)}] {r["id"][:38]:38s} '
              f'{"SOLVED d" + str(res.get("depth")) if res["solved"] else "fail  "} '
              f'{time.time() - t1:5.0f}s | cum solved {solved}/{done} '
              f'| eta {el / done * (len(todo) - done) / 60:.1f} min', flush=True)
        if res['solved']:
            print(f'      proof: {" ; ".join(res["proof"]).replace(chr(10), " ")[:150]}', flush=True)

print(f'\n=== {A.split} slice {A.start}..{A.start + len(todo)}: passed {solved}/{len(todo)} '
      f'= {solved / max(1, len(todo)):.3f}', flush=True)
print(f'lean calls {S.verifier.n_calls} in {S.verifier.lean_seconds:.0f}s '
      f'(avg {S.verifier.lean_seconds / max(1, S.verifier.n_calls):.1f}s/call), '
      f'total wall {(time.time() - t_start) / 60:.1f} min', flush=True)
print(f'results -> {out_path}', flush=True)