import argparse, json
from lamm.data import stream_hf, batched
from lamm.verify import ProofAuditor
from lamm.data import format_sft
from lamm.metrics import summarize

def main():
    ap=argparse.ArgumentParser()
    ap.add_argument("--dataset",default="nvidia/Nemotron-SFT-Math-v4")
    ap.add_argument("--split",default="train")
    ap.add_argument("--source",default="StackExchange-Math")
    ap.add_argument("--batch-size",type=int,default=16)
    ap.add_argument("--batches",type=int,default=10)
    args=ap.parse_args()

    auditor=ProofAuditor()
    audits=[]
    for bi,batch in enumerate(batched(stream_hf(args.dataset,args.split,args.source),args.batch_size),1):
        print(f"BATCH {bi}: {len(batch)}")
        for rec in batch:
            assistant = ""
            for m in reversed(rec.messages):
                if isinstance(m,dict) and m.get("role")=="assistant":
                    assistant=str(m.get("content") or m.get("reasoning_content") or "")
                    break
            if not assistant:
                continue
            import re
            steps=[x.strip() for x in re.split(r"(?i)(?=\\bSTEP\\s*\\d+\\b)|\\n\\s*\\n",assistant) if x.strip()]
            if len(steps)==1:
                steps=[assistant]
            a=auditor.audit(steps,rec.expected_answer)
            audits.append(a)
        if bi>=args.batches: break
    print(json.dumps(summarize(audits),indent=2))

if __name__=="__main__":
    main()