"""Self-tuning loop: an optimizer model rewrites the worker's skill from its failures on the tune set.

Usage: autotune.py <lang> <rep> [--gate] [--rounds 8]
- Items are split at random (seed = rep) into tune / holdout / test.
- Ungated: every proposal is accepted; we only record tune, holdout and test scores each round.
- Gated (3.11): a proposal is accepted only if the holdout score rises by at least MARGIN and it breaks no more
  holdout items than it fixes (--gate-strict, the first version: no drop at all and no item may break); otherwise we keep the current skill and tell the optimizer the proposal was rejected
  (without showing any holdout item). Stop after 3 rejections in a row.
The optimizer never sees holdout or test items. Test is only used for the final report.
"""
import json, random, sys
from pathlib import Path

sys.path.insert(0, "<chatgpt-oauth>/src")
from chat_gpt_oauth.client import request_result  # noqa: E402
from lib import HERE, aggregate, run_skill  # noqa: E402

lang, rep = sys.argv[1], int(sys.argv[2])
GATE = "--gate" in sys.argv or "--gate-strict" in sys.argv
STRICT = "--gate-strict" in sys.argv
AVG = "--gate-avg" in sys.argv
if AVG: GATE = True
MARGIN = 0.01  # about the run-to-run noise of a stable skill on the 60-item holdout
ROUNDS = int(sys.argv[sys.argv.index("--rounds") + 1]) if "--rounds" in sys.argv else 8
N_TUNE, N_HOLD = 24, 60
OPT_MODEL = "gpt-6.1-sol"

items = json.loads((HERE / f"items_{lang}.json").read_text(encoding="utf-8"))
names = (HERE / f"names_{lang}.md").read_text(encoding="utf-8")
rng = random.Random(rep)
order = items[:]; rng.shuffle(order)
split = {"tune": order[:N_TUNE], "holdout": order[N_TUNE:N_TUNE + N_HOLD], "test": order[N_TUNE + N_HOLD:]}
out_dir = HERE / "runs" / f"{lang}_{('gatedavg' if AVG else 'gatedstrict' if STRICT else 'gated') if GATE else 'ungated'}_{rep}"
out_dir.mkdir(parents=True, exist_ok=True)
(out_dir / "split.json").write_text(json.dumps({k: [i["id"] for i in v] for k, v in split.items()}, indent=1), encoding="utf-8")
META = (HERE / f"meta_prompt_{lang}.md").read_text(encoding="utf-8")


def evaluate(skill):
    res = {k: run_skill(skill, names, v, lang) for k, v in split.items()}
    if AVG:  # second, independent holdout run; the gate compares the mean of the two
        res["holdout_rerun"] = run_skill(skill, names, split["holdout"], lang)
    return res


def hold_score(res):
    a = aggregate(res["holdout"])["score"]
    return (a + aggregate(res["holdout_rerun"])["score"]) / 2 if AVG else a


def failure_report(tune_res):
    by_id = {i["id"]: i for i in split["tune"]}
    lines = []
    for r in tune_res:
        if r["misses"] or r["extras"]:
            it = by_id[r["id"]]
            lines.append(json.dumps({"id": r["id"], "raw": it["raw"], "output": r["output"], "expected": it["truth"],
                                     "missed": r["misses"], "wrongly_added": r["extras"]}, ensure_ascii=False))
    return "\n".join(lines) if lines else "(no failures)"


def propose(skill, tune_res, history_note):
    agg = aggregate(tune_res)
    prompt = (META.replace("{skill}", skill).replace("{names}", names)
              .replace("{score}", f"{agg['score']}（写对 {agg['correct']}/{agg['total']}，乱改 {agg['inserts']}）" if lang == "zh"
                       else f"{agg['score']} (correct {agg['correct']}/{agg['total']}, wrongly added {agg['inserts']})")
              .replace("{failures}", failure_report(tune_res)).replace("{history}", history_note))
    ans = request_result(prompt, model=OPT_MODEL, reasoning_effort="high", timeout=900)["answer"].strip()
    if ans.startswith("```"):
        ans = ans.strip("`").split("\n", 1)[1]
    return ans.strip()


skill = (HERE / f"skill_v0_{lang}.md").read_text(encoding="utf-8").strip()
res = evaluate(skill)
log = [{"round": 0, "accepted": True, "skill_chars": len(skill), **{k: aggregate(v) for k, v in res.items()}}]
if AVG: log[0]["holdout_mean"] = round(hold_score(res), 4)
(out_dir / "skill_r00.md").write_text(skill + "\n", encoding="utf-8")
print(json.dumps(log[-1], ensure_ascii=False), flush=True)
current, cur_res, rejects, note = skill, res, 0, ""
for rnd in range(1, ROUNDS + 1):
    cand = propose(current, cur_res["tune"], note)
    (out_dir / f"skill_r{rnd:02d}.md").write_text(cand + "\n", encoding="utf-8")
    cres = evaluate(cand)
    entry = {"round": rnd, "skill_chars": len(cand), **{k: aggregate(v) for k, v in cres.items()}}
    if GATE:
        ok_before = {r["id"]: not r["misses"] and not r["extras"] for r in cur_res["holdout"]}
        ok_after = {r["id"]: not r["misses"] and not r["extras"] for r in cres["holdout"]}
        broken = sorted(i for i in ok_before if ok_before[i] and not ok_after[i])
        fixed = sorted(i for i in ok_before if not ok_before[i] and ok_after[i])
        if STRICT:  # first version: no drop at all, and no holdout item may break
            ok = entry["holdout"]["score"] >= aggregate(cur_res["holdout"])["score"] and not broken
        else:       # must beat the noise margin, and break no more items than it fixes
            ok = hold_score(cres) >= hold_score(cur_res) + MARGIN and len(broken) <= len(fixed)
            entry["holdout_mean"] = round(hold_score(cres), 4)
        entry.update(accepted=ok, regressions=broken, fixed=fixed)
        if ok:
            current, cur_res, rejects = cand, cres, 0
            note = ""
        else:
            rejects += 1
            note = ("上一轮你提出的修改没有通过门禁：在你看不到的留出集上，它让分数下降或者把原来对的条目改错了。"
                    "下面附的仍是当前生效版本的失败情况，请换一个思路修改。" if lang == "zh" else
                    "Your previous proposal was rejected by the gate: on a holdout set you cannot see, it lowered the score or broke "
                    "items that used to be right. The failures below are still for the version currently in force; try a different approach.")
    else:
        entry["accepted"] = True
        current, cur_res = cand, cres
    log.append(entry)
    (out_dir / f"results_r{rnd:02d}.json").write_text(json.dumps(cres, ensure_ascii=False), encoding="utf-8")
    (out_dir / "log.json").write_text(json.dumps(log, ensure_ascii=False, indent=1), encoding="utf-8")
    print(json.dumps({k: v for k, v in entry.items() if k != "regressions"}, ensure_ascii=False), flush=True)
    if GATE and rejects >= 3:
        break
(out_dir / "final_skill.md").write_text(current + "\n", encoding="utf-8")
(out_dir / "log.json").write_text(json.dumps(log, ensure_ascii=False, indent=1), encoding="utf-8")