"""5.5 开不开思考：同一个模型、同一批带答案的题，关掉和打开思考各跑几遍，
比较准确率、延迟、输出 token 和花费。

用法：
    export AI_BUILDER_TOKEN=sk_...        # 你自己的 AI Builder Space key
    python run.py                          # 默认 grok-4-fast，16 道示例题，每种模式 3 遍
    python run.py --model gpt-5 --repeats 2
    python run.py --tasks my_tasks.jsonl   # 换成你自己的题

题目文件每行一个 JSON：{"id": "...", "question": "...", "answer": ["可接受的答案1", "答案2"]}
"""
import argparse
import concurrent.futures as cf
import json
import os
import re
import statistics
import time
from datetime import datetime
from fractions import Fraction
from pathlib import Path

from openai import OpenAI

BASE_URL = os.getenv('AI_BUILDER_SPACE_BASE_URL', 'https://space.ai-builders.com/backend/v1')

# 美元 / 每百万 token（输入, 输出）。AI Builder Space 2026-10 的价目，换模型或价格变了就改这里。
PRICES = {
    'grok-4-fast': (1.25, 2.50),
    'gpt-5': (4.00, 20.00),
    'gpt-6': (2.00, 10.00),
    'deepseek': (0.15, 0.60),
}

# 两种模式只差 reasoning_effort 这一个参数，模型、题目、提示都相同。
# grok-4-fast 和 gpt-5 都接受 none（不思考）和 low / medium / high。
MODES = {'off': 'none', 'on': 'high'}

SYSTEM = '回答下面的问题。可以先写推理过程，但最后必须单独一行写“答案：<最终答案>”，答案里不要加单位。'


def normalize(s):
    s = s.strip().lower()
    s = re.sub(r'\\d?frac\{([^}]*)\}\{([^}]*)\}', r'\1/\2', s)        # LaTeX 分数
    s = re.sub(r'[*`$\\()\[\]{}]', '', s)
    return s.strip().strip('。.，,：: ').replace(' ', '')


def as_number(s):
    try:
        return Fraction(s.replace(',', ''))
    except (ValueError, ZeroDivisionError):
        return None


def extract_answer(reply):
    """取最后一个“答案：”后面的内容；模型没按格式写时，退回到最后一行。"""
    hits = re.findall(r'答案\s*[:：]\s*(.+)', reply or '')
    if hits:
        return hits[-1]
    lines = [l for l in (reply or '').strip().splitlines() if l.strip()]
    return lines[-1] if lines else ''


def is_correct(reply, answers):
    got = normalize(extract_answer(reply))
    if not got:
        return False
    for a in answers:
        want = normalize(a)
        if got == want:
            return True
        w = as_number(want)
        nums = re.findall(r'-?\d+(?:\.\d+)?(?:/\d+)?', got.replace(',', ''))   # “130度”“0.05美元”只取数字
        if w is not None and nums and abs(float(as_number(nums[0])) - float(w)) < 1e-6:
            return True
    return False


def ask(client, model, effort, question):
    t0 = time.time()
    r = client.chat.completions.create(
        model=model,
        messages=[{'role': 'system', 'content': SYSTEM}, {'role': 'user', 'content': question}],
        max_tokens=16000,             # 思考会消耗输出 token，给足预算，否则可能返回空
        reasoning_effort=effort,
    )
    return r.choices[0].message.content or '', r.usage, time.time() - t0


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--model', default='grok-4-fast')
    ap.add_argument('--tasks', default=str(Path(__file__).with_name('tasks.jsonl')))
    ap.add_argument('--repeats', type=int, default=3, help='每种模式把整套题跑几遍')
    ap.add_argument('--workers', type=int, default=4, help='同时发几个请求')
    ap.add_argument('--out', default='results')
    args = ap.parse_args()

    token = os.getenv('AI_BUILDER_TOKEN')
    if not token:
        raise SystemExit('请先设置环境变量 AI_BUILDER_TOKEN（你的 AI Builder Space key）')
    client = OpenAI(base_url=BASE_URL, api_key=token, timeout=600)
    price_in, price_out = PRICES.get(args.model, (0, 0))
    tasks = [json.loads(l) for l in Path(args.tasks).read_text(encoding='utf-8').splitlines() if l.strip()]

    jobs = [(mode, rep, t) for mode in MODES for rep in range(args.repeats) for t in tasks]
    rows = []

    def run(job):
        mode, rep, t = job
        try:
            reply, usage, sec = ask(client, args.model, MODES[mode], t['question'])
        except Exception as e:                                   # 失败也记下来，算答错
            return dict(mode=mode, rep=rep, id=t['id'], ok=False, sec=None, error=str(e)[:200],
                        prompt_tokens=0, completion_tokens=0, cost=0.0, reply='')
        cost = (usage.prompt_tokens * price_in + usage.completion_tokens * price_out) / 1e6
        return dict(mode=mode, rep=rep, id=t['id'], ok=is_correct(reply, t['answer']), sec=round(sec, 2),
                    prompt_tokens=usage.prompt_tokens, completion_tokens=usage.completion_tokens,
                    cost=cost, reply=reply[-300:])

    print(f'模型 {args.model}，{len(tasks)} 道题，每种模式 {args.repeats} 遍，共 {len(jobs)} 次调用……')
    with cf.ThreadPoolExecutor(args.workers) as ex:
        for i, row in enumerate(ex.map(run, jobs), 1):
            rows.append(row)
            if i % 10 == 0:
                print(f'  {i}/{len(jobs)}')

    print(f'\n## 汇总（{args.model}，思考关 = reasoning_effort {MODES["off"]}，思考开 = {MODES["on"]}）\n')
    print('| 模式 | 准确率（各遍） | 延迟中位数 | 平均输出 token | 整套题一遍的花费 |')
    print('|---|---|---|---|---|')
    summary = {}
    for mode in MODES:
        rs = [r for r in rows if r['mode'] == mode]
        per_rep = [sum(r['ok'] for r in rs if r['rep'] == k) for k in range(args.repeats)]
        secs = [r['sec'] for r in rs if r['sec'] is not None]
        out_tok = statistics.mean(r['completion_tokens'] for r in rs)
        cost = sum(r['cost'] for r in rs) / args.repeats
        summary[mode] = dict(correct_per_repeat=per_rep, n_tasks=len(tasks), median_sec=statistics.median(secs),
                             mean_completion_tokens=round(out_tok, 1), cost_per_pass_usd=cost,
                             errors=sum(1 for r in rs if r.get('error')))
        label = '思考关' if mode == 'off' else '思考开'
        print(f'| {label} | {statistics.mean(per_rep):.1f}/{len(tasks)}（{", ".join(map(str, per_rep))}） '
              f'| {statistics.median(secs):.1f} 秒 | {out_tok:.0f} | ${cost:.4f} |')

    print('\n## 每道题答对几遍（关 → 开）\n')
    for t in tasks:
        off = sum(r['ok'] for r in rows if r['id'] == t['id'] and r['mode'] == 'off')
        on = sum(r['ok'] for r in rows if r['id'] == t['id'] and r['mode'] == 'on')
        mark = '  ← 思考有用' if on > off else ('  ← 思考反而更差' if on < off else '')
        print(f'- {t["id"]}: {off}/{args.repeats} → {on}/{args.repeats}{mark}')

    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    stamp = datetime.now().strftime('%Y%m%d_%H%M')
    path = out / f'{args.model}_{stamp}.json'
    path.write_text(json.dumps(dict(model=args.model, modes=MODES, repeats=args.repeats, summary=summary, rows=rows),
                               ensure_ascii=False, indent=1), encoding='utf-8')
    print(f'\n逐条结果（含模型最后的回复）已存到 {path}')


if __name__ == '__main__':
    main()