"""段数と成功率の実験。10件の注文の請求額を、前の累計に足し込んで最終累計を出す仕事を、
split（1段あたりに処理する注文数のリスト）で何段に分けるか変えて測る。
各段は独立の `claude -p`（ツールなし）。前段のモデル出力（累計）が次段の入力になる。
使い方: python3 chain.py <config名> <本数> <出力jsonl>"""
import json, math, os, random, subprocess, sys, tempfile, time
from collections import Counter
from concurrent.futures import ThreadPoolExecutor

NAMES = ["りんご","みかん","牛乳","食パン","卵","バナナ","豆腐","納豆","ヨーグルト","チーズ","ハム","鶏肉","豚肉","トマト","キャベツ","人参","玉ねぎ","じゃがいも"]
SYS = "あなたは正確に作業するアシスタントです。指示された出力形式だけを返してください。"
RULE = ("請求額の求め方: 各行を解析し、qty*price を小計とする。qty>=5 の行だけ小計を0.9倍して小数点以下切り捨て。"
        "全行の合計に消費税10%(切り捨て)を加えた税込額 gross を出し、gross が5000以上なら送料0、未満なら送料600円を足したものが請求額。")

# config: model, 1注文あたりの行数、分割、1段あたりの投票数
CONFIGS = {
  # 実験B: 各段の難しさを一定（1段=注文1件）。途中経過から1/3/6/10段を同時に得る
  "B_haiku_8":   dict(model="haiku",  lines=8, split=[1]*10, votes=1),
  "B_sonnet_8":  dict(model="sonnet", lines=8, split=[1]*10, votes=1),
  "B_haiku_3":   dict(model="haiku",  lines=3, split=[1]*10, votes=1),
  "B_haiku_8_v3":dict(model="haiku",  lines=8, split=[1]*10, votes=3),
  # 実験C: AIは注文テキストの読み取り（数量・単価の抽出）だけ。割引・税・送料・累計はコードで計算
  "C_haiku_8_parse": dict(model="haiku", lines=8, split=[1]*10, votes=1, mode="parse"),
  # 実験A: 同じ仕事（10件・各6行）を1/3/6/10段に分ける
  "A_haiku_1":   dict(model="haiku",  lines=6, split=[10], votes=1),
  "A_haiku_3":   dict(model="haiku",  lines=6, split=[4,3,3], votes=1),
  "A_haiku_6":   dict(model="haiku",  lines=6, split=[2,2,2,2,1,1], votes=1),
  "A_haiku_10":  dict(model="haiku",  lines=6, split=[1]*10, votes=1),
}

def make_order(seed, lines):
    r = random.Random(seed); names = r.sample(NAMES, lines); out = []; bill = []
    for nm in names:
        q = r.randint(1, 9); p = r.randint(48, 980)
        out.append(r.choice(["{n} x{q} @{p}円","{n} {q}個 単価{p}","{n} / 数量{q} / ¥{p}","{q}点 {n} ({p}円)"]).format(n=nm, q=q, p=p))
        bill.append(math.floor(q*p*0.9) if q >= 5 else q*p)
    total = sum(bill); g = total + total//10
    return "\n".join(out), g + (0 if g >= 5000 else 600)

def call(prompt, model):
    cmd = ["claude","-p",prompt,"--model",model,"--tools","","--output-format","json",
           "--setting-sources","","--strict-mcp-config","--disable-slash-commands",
           "--no-session-persistence","--system-prompt",SYS]
    t0 = time.time()
    try:
        p = subprocess.run(cmd, cwd=tempfile.mkdtemp(prefix="stepexp_"), capture_output=True, text=True, timeout=300, stdin=subprocess.DEVNULL)
        res = json.loads(p.stdout)
    except Exception as e:
        return {"result": "", "cost": 0, "wall": time.time()-t0, "models": [], "err": str(e)[:80]}
    return {"result": res.get("result",""), "cost": res.get("total_cost_usd",0), "wall": time.time()-t0,
            "models": list((res.get("modelUsage") or {}).keys()), "err": None if not res.get("is_error") else "is_error"}

def parse(s):
    a, b = s.find("{"), s.rfind("}")
    try: v = json.loads(s[a:b+1])["running"]; return v if isinstance(v, int) else None
    except Exception: return None

def bill(items):
    sub = [math.floor(q*p*0.9) if q >= 5 else q*p for q, p in items]; total = sum(sub); g = total + total//10
    return g + (0 if g >= 5000 else 600)

def run_parse_job(name, seed):
    c = CONFIGS[name]; run = truth = 0; steps = []; cost = wall = 0; models = set()
    for i in range(10):
        text, ans = make_order(seed*1000 + i, c["lines"]); truth += ans
        r = call("次の注文テキストを解析し、JSON配列 [{\"qty\":整数,\"price\":整数(単価円)}] だけを返せ。行ごとに1要素。\n" + text, c["model"])
        cost += r["cost"]; wall += r["wall"]; models.update(r["models"])
        try:
            res = r["result"]; items = json.loads(res[res.find("["):res.rfind("]")+1]); got_bill = bill([(int(x["qty"]), int(x["price"])) for x in items])
        except Exception: got_bill = None
        prev = run; run = run + got_bill if got_bill is not None else run
        steps.append({"orders_done": i+1, "prev": prev, "got": run, "local_ok": got_bill == ans, "cum_ok": run == truth, "votes": [run]})
    return {"config": name, "seed": seed, "steps": steps, "cost": cost, "wall": wall, "calls": 10, "models": sorted(models)}

def run_job(args):
    name, seed = args; c = CONFIGS[name]
    if c.get("mode") == "parse":
        return run_parse_job(name, seed)
    run = 0; truth = 0; idx = 0; steps = []; cost = wall = 0; models = set(); calls = 0
    for size in c["split"]:
        texts = []; add = 0
        for j in range(size):
            t, ans = make_order(seed*1000 + idx, c["lines"]); idx += 1; texts.append(t); add += ans
        body = "\n".join(f"注文{j+1}:\n{t}" for j, t in enumerate(texts))
        prompt = (f"次の{size}件の注文それぞれの請求額を求め、前の累計にすべて足せ。{RULE}\n前の累計: {run}\n{body}\n"
                  "出力は {\"running\":整数(前の累計+全注文の請求額)} のJSONだけ。")
        votes = []
        for _ in range(c["votes"]):
            r = call(prompt, c["model"]); calls += 1; cost += r["cost"]; wall += r["wall"]; models.update(r["models"])
            votes.append(parse(r["result"]))
        valid = [v for v in votes if v is not None]
        got = Counter(valid).most_common(1)[0][0] if valid else None
        truth += add
        steps.append({"orders_done": idx, "prev": run, "got": got, "local_ok": got == run + add, "cum_ok": got == truth, "votes": votes})
        run = got if got is not None else 0
    return {"config": name, "seed": seed, "steps": steps, "cost": cost, "wall": wall, "calls": calls, "models": sorted(models)}

if __name__ == "__main__":
    name, n, outp = sys.argv[1], int(sys.argv[2]), sys.argv[3]
    done = set()
    if os.path.exists(outp): done = {json.loads(l)["seed"] for l in open(outp)}
    seeds = [s for s in range(n) if s not in done]
    with ThreadPoolExecutor(int(os.environ.get("EXP_PAR", 8))) as ex, open(outp, "a") as f:
        for row in ex.map(run_job, [(name, s) for s in seeds]):
            f.write(json.dumps(row, ensure_ascii=False) + "\n"); f.flush()