#!/usr/bin/env python3
"""exp15 — pre-reg budget recompute ($0), from live prices (freeze_facts.json).

Judge budget is recomputed FROM trap-N x votes x 2 orders x >=2 judges (prompt discipline: not the
other way round). Hard cap $15 (owner). Per-call predicted-cost gate lives in the run scripts; this
sizes the PACKAGE and picks a judge pair that fits. All assumptions are explicit + conservative so the
freeze can pin them. Self-review: prints per-Q breakdown; asserts total logic.

Cost model (USD/1M tokens from freeze_facts.json):
  arm call  = in_tok*price_in*(1-cache_frac) + in_tok*price_hit*cache_frac + out_tok*price_out
  judge call= in_tok*price_in + out_tok(+reasoning)*price_out           (judges: no prefix cache assumed)
Order of Qs (prompt): Q2b($0) -> Q1+Q3 (S2, carrying) -> Q2a/Q2c -> Q4 -> Q0. Budget likely covers
the carrying block; later Qs run only if budget remains (pre-reg fixes priority, stop is legitimate).
"""
from __future__ import annotations
import json
from pathlib import Path

FACTS = json.load(open("/home/ubuntu/books/gu-zhenren/exp15/freeze_facts.json"))
P = FACTS["prices"]


def price(model, field, tier=0):
    v = P[model][field]
    return v[tier] if isinstance(v, list) else v


def arm_call(model, in_tok, out_tok, cache_frac=0.0):
    pin = price(model, "in"); phit = P[model].get("cache_hit", pin); pout = price(model, "out")
    return (in_tok * pin * (1 - cache_frac) + in_tok * phit * cache_frac + out_tok * pout) / 1e6


def judge_call(model, in_tok, out_tok):
    # reasoning/thinking billed as output for grok/gemini/gpt/kimi -> fold into out_tok by caller
    return (in_tok * price(model, "in") + out_tok * price(model, "out")) / 1e6


# ---------------- ARM (generation) sizing ----------------
# Conservative per-chunk token estimates (zh->ru; injection small; reflow-invariant).
DRAFT_IN, DRAFT_OUT = 2600, 3200      # deepseek-v4-flash: src+small injection in; ru + thinking (billed out)
EDIT_IN, EDIT_OUT = 5200, 2800        # glm-5: src+draft+injection in; ru edits out (thinking off for editor)
CACHE_FRAC = 0.55                     # static prefix (system+glossary) cached on sequential chunks (verified auto-cache)
CHUNKS_PER_CH = 2.2                   # ~2 greedy-1500 chunks per webnovel chapter (S2 est; S1 ~1/节)
S2_CH = 10                            # carrying-slice chapters (8-12)

def arm_cost_per_chapter(draft="deepseek-v4-flash", editor="glm-5"):
    per_chunk = arm_call(draft, DRAFT_IN, DRAFT_OUT, CACHE_FRAC) + arm_call(editor, EDIT_IN, EDIT_OUT, CACHE_FRAC)
    return per_chunk * CHUNKS_PER_CH

# Arms that GENERATE (Q0 do-nothing = $0 export; Q3 modes share one draft, differ in editor-context):
ARMS_GEN = {
    "A0": 1, "A0'": 1,                # double anchor
    "Q1-logical": 1,
    "Q2a-editGLava": 1, "Q2c-wideread": 1,
    "Q3a-src": 1, "Q3a-draft": 1, "Q3a-final": 1, "Q3a'-state": 1, "Q3b-lookahead": 1,
    "Q4a-strongdraft(v4pro)": 1, "Q4b-diff": 1, "Q4c-reflow": 1,
}

# ---------------- JUDGE sizing (recompute from trap-N) ----------------
# Trap set (common denominator, defined in A0). Minimums per prompt D1.
TRAP_N = {"T-ana": 20, "T-ent": 20, "T-tense": 20, "T-cat": 10, "T-ell": 10,
          "T-gen": 10, "T-reg": 10, "T-inv": 10}
# Which Qs judge which trap types, and #arm-contrasts each Q compares against A0 (pairwise A0-vs-arm):
Q_JUDGE = {
    "Q1":    (["T-tense", "T-ana"], 1),      # A0 vs logical
    "Q3a":   (["T-ana", "T-ell"], 3),        # A0 vs {src,draft,final}
    "Q3a'":  (["T-gen", "T-reg"], 1),        # A0 vs +state
    "Q3b":   (["T-cat"], 1),                 # A0 vs +lookahead
    # later Qs (run only if budget remains):
    "Q2c":   (["T-ana", "T-tense"], 1),
    "Q4a":   (["T-inv"], 1),
    "Q4b":   (["T-inv"], 1),
}
CARRYING = {"Q1", "Q3a", "Q3a'", "Q3b"}
VOTES_BASE = 6            # 3 per order x 2 orders
ESCALATE_FRAC = 0.35     # fraction of cells that split -> +4 votes to 10
JUDGES = ["grok-4.3", "mistral-large-latest"]     # 2 cross-family; family-clean vs DeepSeek/GLM arms + kimi miner
JUDGE_IN, JUDGE_OUT = {"grok-4.3": 1500, "mistral-large-latest": 1500}, {"grok-4.3": 900, "mistral-large-latest": 500}
# grok reasoning ON (additive) -> larger out; mistral no reasoning surcharge.

def votes_expected():
    return VOTES_BASE + ESCALATE_FRAC * 4  # ~7.4

def judge_cost_for(qset, judges, jin, jout):
    total = 0.0; breakdown = {}
    for q, (types, contrasts) in Q_JUDGE.items():
        if q not in qset:
            continue
        traps = sum(TRAP_N[t] for t in types)
        cells = traps * contrasts                       # pairwise (A0-vs-arm) trap cells
        calls_per_judge = cells * votes_expected()
        c = 0.0
        for j in judges:
            c += calls_per_judge * judge_call(j, jin[j], jout[j])
        breakdown[q] = (traps, contrasts, cells, round(c, 2))
        total += c
    return total, breakdown

# ---------------- MINER (trap mining, paid) ----------------
# kimi-k2.6 mines candidate cohesion phenomena over the S2 source; verbose reasoning billed as output.
MINER = "kimi-k2.6"
MINER_CALLS = S2_CH            # ~1 mining pass per chapter
MINER_IN, MINER_OUT = 3000, 6000   # chapter source in; candidate list + verbose reasoning out
miner_cost = MINER_CALLS * judge_call(MINER, MINER_IN, MINER_OUT)


def main():
    per_ch = arm_cost_per_chapter()
    arm_total = sum(ARMS_GEN.values()) * S2_CH * per_ch
    carrying_arm = 6 * S2_CH * per_ch   # A0,A0',Q1,Q3a(counts as 1 draft+3 editor-ctx ~ approx 3),Q3a',Q3b ~6 arm-equivs
    j_all, jb = judge_cost_for(set(Q_JUDGE), JUDGES, JUDGE_IN, JUDGE_OUT)
    j_carry, jbc = judge_cost_for(CARRYING, JUDGES, JUDGE_IN, JUDGE_OUT)

    print("=== exp15 budget recompute (live prices, 2026-07-17) ===")
    print(f"arm cost / chapter (draft deepseek-v4-flash + editor glm-5, cache_frac={CACHE_FRAC}): ${per_ch:.4f}")
    print(f"  draft/chunk ${arm_call('deepseek-v4-flash',DRAFT_IN,DRAFT_OUT,CACHE_FRAC):.4f}  editor/chunk ${arm_call('glm-5',EDIT_IN,EDIT_OUT,CACHE_FRAC):.4f}  chunks/ch {CHUNKS_PER_CH}")
    print(f"trap miner ({MINER}, {MINER_CALLS} calls): ${miner_cost:.2f}")
    print(f"expected votes/cell: {votes_expected():.1f}  judges: {JUDGES}")
    print("\n-- judge cost per Q (traps, contrasts, cells, $) --")
    for q, v in jb.items():
        tag = "CARRY" if q in CARRYING else "later"
        print(f"  [{tag}] {q}: traps={v[0]} contrasts={v[1]} cells={v[2]} -> ${v[3]}")
    print(f"\nCARRYING block: arms ${carrying_arm:.2f} + judges ${j_carry:.2f} + miner ${miner_cost:.2f} = ${carrying_arm+j_carry+miner_cost:.2f}")
    print(f"FULL package (all listed Qs): arms ${arm_total:.2f} + judges ${j_all:.2f} + miner ${miner_cost:.2f} = ${arm_total+j_all+miner_cost:.2f}")
    print(f"HARD CAP: $15.00")
    carry_total = carrying_arm + j_carry + miner_cost
    print(f"\nVERDICT: carrying block ${carry_total:.2f} {'FITS' if carry_total<=15 else 'EXCEEDS'} $15; "
          f"later Qs (Q2c/Q4) run only if budget remains after carrying block + per-call gate.")
    # sensitivity: apex secondary judge (gemini mandatory thinking billed as output)
    apex_judges = ["grok-4.3", "gemini-3.1-pro-preview"]
    apex_in = {**JUDGE_IN, "gemini-3.1-pro-preview": 1500}
    apex_out = {**JUDGE_OUT, "gemini-3.1-pro-preview": 2500}
    jg, _ = judge_cost_for(CARRYING, apex_judges, apex_in, apex_out)
    print(f"\n[sensitivity] if secondary judge = gemini-3.1-pro (apex, thinking billed out): carrying judges ${jg:.2f} "
          f"-> block ${carrying_arm+jg+miner_cost:.2f} ({'FITS' if carrying_arm+jg+miner_cost<=15 else 'EXCEEDS'} $15)")


if __name__ == "__main__":
    main()
