#!/usr/bin/env python3
"""experiments/08 — COGS v2 on the ratified Phase-1 stack (draft=DeepSeek-v4-flash,
editor=grok-4.3, premium apex=Gemini-3.1-Pro). Supersedes exp01's Sonnet-editor numbers.

WHY v2: the editor changed from a cheap model (GLM/DeepSeek, ~$0.28/M out) to grok-4.3
($2.50/M out ≈ 9× draft output). Because the editor output is the same volume as the
draft (~r×B tokens) but ~9× the price, the EDITOR now dominates standard COGS — a single
blended M_out multiplier (exp01) can no longer be used; each stage is priced by its own model.

PRICES — official docs, dated (fetched + cross-verified 2026-07-05; NOT from memory):
  deepseek-v4-flash  in $0.14  cache-hit $0.0028  out $0.28   (reasoning billed as output; no batch)
                     src: api-docs.deepseek.com/quick_start/pricing
  grok-4.3           in $1.25  cached-in $0.20    out $2.50   (reasoning additive→output; no batch)
                     src: docs.x.ai/developers/models/grok-4.3
  gemini-3.1-pro     in $2.00(<=200k)/$4.00(>200k)  out $12.00/$18.00 (INCLUDES forced thinking)
                     cached-in $0.20/$0.40  BATCH -50%   src: ai.google.dev/gemini-api/docs/pricing (upd 2026-06-30)

TOKEN MODEL (per book of B source-tokens, output ratio r=out/in from exp01 parallel pairs):
  Draft (bilingual translator):  in = (1 + h_draft)·B   out = r·B      [+ small deepseek reasoning, in noise]
  Editor (MONOLINGUAL, D1):      in = (r + h_edit)·B     out = r·B      [thinking-OFF, ledger clean]
  Apex  (Gemini bilingual final, forced thinking): in = (1.5 + r)·B  out = r·B·(1 + rf_gem)
  Judge (Gemini, 20% sample, batch): in = 0.2·(1+r)·B   out = 0.1·r·B·(1 + rf_gem)
  h_draft=1.0 (source + system + selective glossary + STM overhead ≈ B);
  h_edit=0.5  (monolingual editor sees the draft r·B + a small prompt/glossary, no source, no big STM);
  rf_gem = Gemini thinking-as-output factor (central 1.0; sensitivity 0.5–2.0).

BOOKS (B, r from exp01-token-calibration, direct-keys contour, EU SaaS):
  ja→ru ранобэ (том)      B=113_000    r=1.29
  zh→ru вебновелла 500гл  B=1_290_000  r=1.88
  en→ru роман             B=131_000    r=1.53
"""
from __future__ import annotations
import json
from pathlib import Path

P = {  # USD per token (per-1M / 1e6)
    "ds_in": 0.14e-6, "ds_hit": 0.0028e-6, "ds_out": 0.28e-6,
    "grok_in": 1.25e-6, "grok_hit": 0.20e-6, "grok_out": 2.50e-6,
    "gem_in": 2.00e-6, "gem_out": 12.00e-6,               # <=200k tier (chunks are ~1.5k tok → always <=200k)
    "gem_in_b": 1.00e-6, "gem_out_b": 6.00e-6,            # batch -50%
}
H_DRAFT, H_EDIT = 1.0, 0.5
BOOKS = {
    "ja→ru ранобэ (том)":      {"B": 113_000,   "r": 1.29},
    "zh→ru вебновелла 500гл":  {"B": 1_290_000, "r": 1.88},
    "en→ru роман":             {"B": 131_000,   "r": 1.53},
}

def draft_cost(B, r, cache_frac=0.0):
    """DeepSeek draft. cache_frac = fraction of INPUT served from auto-cache (stable prefix)."""
    tin = (1 + H_DRAFT) * B
    hit = tin * cache_frac
    miss = tin - hit
    c_in = miss * P["ds_in"] + hit * P["ds_hit"]
    c_out = r * B * P["ds_out"]
    return c_in + c_out, c_in, c_out

# Editor = grok-4.3 + explicit reasoning_effort:"none" → reasoning=0 (grok default is `low`,
# billed additively as output). reasoning_factor=0 = editor default (off); 0.25 = reasoning left
# ON, +25% output. Off-switch mechanism/probe → docs/experiments/00-provider-quirks.md §grok.
EDITOR_REASONING_FACTOR = 0.0

def editor_cost(B, r, cache_frac=0.0, reasoning_factor=0.0):
    """Monolingual editor (grok-4.3 + reasoning_effort:"none"). reasoning_factor=0 → reasoning off
    (default); 0.25 → reasoning left ON (grok default `low`), billed additively as output."""
    tin = (r + H_EDIT) * B
    hit = tin * cache_frac
    miss = tin - hit
    c_in = miss * P["grok_in"] + hit * P["grok_hit"]
    c_out = r * B * (1 + reasoning_factor) * P["grok_out"]
    return c_in + c_out, c_in, c_out

def apex_cost(B, r, rf):
    """Gemini-3.1-Pro final bilingual pass; forced thinking → output×(1+rf)."""
    c_in = (1.5 + r) * B * P["gem_in"]
    c_out = r * B * (1 + rf) * P["gem_out"]
    return c_in + c_out

def judge_cost(B, r, rf):
    """Gemini judge on 20% sample, batch -50%."""
    c_in = 0.2 * (1 + r) * B * P["gem_in_b"]
    c_out = 0.1 * r * B * (1 + rf) * P["gem_out_b"]
    return c_in + c_out

def money(x): return f"${x:,.3f}" if x < 10 else f"${x:,.2f}"

def main():
    out = {"prices_asof": "2026-07-05", "h_draft": H_DRAFT, "h_edit": H_EDIT, "books": {}}
    for name, bk in BOOKS.items():
        B, r = bk["B"], bk["r"]
        d, d_in, d_out = draft_cost(B, r)
        d_c, dc_in, dc_out = draft_cost(B, r, cache_frac=0.5)  # 50% prefix cache-hit sensitivity
        e, e_in, e_out = editor_cost(B, r)                          # non-reasoning editor (recommended)
        e43, _, _ = editor_cost(B, r, reasoning_factor=0.25)        # grok-4.3 reasoning kept
        std = d + e
        std_grok43 = d + e43
        std_cached = d_c + e
        premium_by_rf = {}
        for rf in (0.5, 1.0, 2.0):
            ap = apex_cost(B, r, rf)
            jg = judge_cost(B, r, rf)
            premium_by_rf[f"rf={rf}"] = {"apex": ap, "judge": jg, "premium_total": std + ap + jg}
        # Selective premium (realistic for large works): apex re-touches only a FRACTION of
        # chunks (flagged/hard); the rest ship at standard. rf=1.0 central.
        ap_full = apex_cost(B, r, 1.0)
        jg_full = judge_cost(B, r, 1.0)
        selective = {f"apex_{int(f*100)}pct": std + f * ap_full + jg_full for f in (0.10, 0.15, 0.25)}
        rec = {
            "B": B, "r": r,
            "draft": {"total": d, "in": d_in, "out": d_out},
            "editor": {"total": e, "in": e_in, "out": e_out},
            "standard_mix": std,
            "standard_mix_grok43_reasoning": std_grok43,
            "standard_mix_50pct_cache": std_cached,
            "editor_share_of_standard": round(e / std, 3),
            "premium_mix": premium_by_rf,
            "selective_premium_rf1": selective,
        }
        out["books"][name] = rec
        print(f"\n=== {name}   (B={B:,} src-tok, r={r}) ===")
        print(f"  draft  (deepseek): {money(d)}   [in {money(d_in)} + out {money(d_out)}]")
        print(f"  editor (grok-4.20-nonreason): {money(e)}   [in {money(e_in)} + out {money(e_out)}]")
        print(f"  STANDARD-MIX (draft+editor):           {money(std)}   (editor = {e/std:.0%} of it)")
        print(f"     if editor = grok-4.3 (reasoning):   {money(std_grok43)}   (+{(std_grok43/std-1)*100:.0f}%)")
        print(f"     with 50% input cache-hit on draft:  {money(std_cached)}")
        for rf, pv in premium_by_rf.items():
            print(f"  PREMIUM full apex (+judge, {rf:>7}): {money(pv['premium_total'])}"
                  f"   [apex {money(pv['apex'])} + judge {money(pv['judge'])}]")
        print(f"  SELECTIVE premium (apex on X% of chunks, rf=1.0):"
              f"  10%={money(selective['apex_10pct'])}  15%={money(selective['apex_15pct'])}  25%={money(selective['apex_25pct'])}")
    Path(__file__).resolve().parent.joinpath("data", "cost_model_v2.json").write_text(
        json.dumps(out, ensure_ascii=False, indent=2))
    print("\nsaved → eval/data/cost_model_v2.json")

if __name__ == "__main__":
    main()