"""Blind re-solve of a deterministic sample of accepted tasks by Qwen3.5-397B-A17B (sampling only)."""
import argparse, asyncio, json, os, sys, time
import common, tinker
from tinker_cookbook import renderers
from tinker_cookbook.tokenizer_utils import get_tokenizer
from tinker_cookbook.recipes.math_rl.math_grading import extract_boxed
sys.path.insert(0, os.path.expanduser("~/repos/research-exchange"))
from research_exchange import generation_audit as ga

SUFFIX = ("\n\nSolve this problem step by step. The answer is a single number. "
          "Put only the final number (no units, no commas) inside \\boxed{}.")


def norm(ans: str) -> str:
    return ans.replace(",", "").replace("$", "").replace("\\!", "").strip()


async def main(a):
    tasks = [json.loads(l) for l in open(a.tasks)]
    evaluation = [json.loads(l) for l in open(a.eval_file)]
    sample = ga.select_sample(tasks, a.sample, a.seed)
    sc = common.service()
    renderer = renderers.get_renderer(a.renderer, tokenizer=get_tokenizer(a.solver))
    sampler = sc.create_sampling_client(base_model=a.solver)
    sem = asyncio.Semaphore(64)
    t0 = time.time()

    async def solve(t):
        prompt = renderer.build_generation_prompt([{"role": "user", "content": t["prompt"] + SUFFIX}])
        async with sem:
            r = await sampler.sample_async(prompt, num_samples=1, sampling_params=tinker.SamplingParams(
                max_tokens=a.max_tokens, temperature=0.0, stop=renderer.get_stop_sequences()))
        msg, _ = renderer.parse_response(r.sequences[0].tokens)
        text = renderers.get_text_content(msg)
        try:
            ans = norm(extract_boxed(text))
        except ValueError:
            ans = ""
        return {"item_id": t["item_id"], "independent_answer": ans, "resolver_id": f"{a.solver} renderer={a.renderer} greedy max_tokens={a.max_tokens}",
                "method": "independent_resolve", "reasoning": text, "blind_to_label": True}

    samples = await asyncio.gather(*[solve(t) for t in sample])
    secs = time.time() - t0
    by_id = {t["item_id"]: t for t in tasks}
    disagreements = []
    for s in samples:
        ag = ga.answer_agreement(by_id[s["item_id"]]["answer"], s["independent_answer"], "numeric_exact")
        if ag is not True:
            disagreements.append({"item_id": s["item_id"], "label": by_id[s["item_id"]]["answer"],
                                  "independent_answer": s["independent_answer"], "agreement": ag})
    accepted = [t for t in tasks if t["disposition"] == "accepted"]
    overlap = ga.evaluation_overlap(accepted, evaluation)
    counts = {"generated": len(tasks), "accepted": len(accepted),
              "rejected": sum(t["disposition"] == "rejected" for t in tasks),
              "execution_errors": sum(t["disposition"] == "execution_error" for t in tasks)}
    report = {"version": 1, "stage": "before_training", "dataset_file": os.path.basename(a.tasks),
              "evaluation_file": os.path.basename(a.eval_file), "audit_file": os.path.basename(a.out),
              "sample_seed": a.seed, "counts": counts, "samples": samples,
              "label_error_count": sum(d["agreement"] is False for d in disagreements),
              "overlap_count": overlap["overlap_count"]}
    v = ga.validate(report, {"required": True, "answer_comparison": "numeric_exact", "min_sample_size": a.sample,
                             "max_label_errors": 0, "max_eval_overlaps": 0},
                    tasks=tasks, evaluation=evaluation, require_comparison=True)
    json.dump(report, open(a.out, "w"), indent=1)
    json.dump({"validation": {k: v[k] for k in v if k != "overlap"}, "overlap": {k: v["overlap"][k] for k in v["overlap"] if k != "matches"} if "overlap" in v else None,
               "disagreements": disagreements, "seconds": secs, "n": len(samples)},
              open(a.out + ".validation.json", "w"), indent=1)
    print(json.dumps({"status": v["status"], "reasons": v["reasons"], "disagreements": disagreements,
                      "seconds": round(secs, 1), "duplicates": v.get("overlap", {}).get("duplicate_count")}, indent=1))


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--tasks", default="/tmp/rx-selfgen/tasks.jsonl")
    p.add_argument("--eval_file", default="/tmp/rx-selfgen/gsm8k_test_eval.jsonl")
    p.add_argument("--solver", default="Qwen/Qwen3.5-397B-A17B")
    p.add_argument("--renderer", default="qwen3_5")
    p.add_argument("--max_tokens", type=int, default=8192)
    p.add_argument("--sample", type=int, default=100)
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--out", default="/tmp/rx-selfgen/audit.json")
    asyncio.run(main(p.parse_args()))