"""5.5 Turning thinking on and off: run the same model on the same labeled tasks
several times with thinking off and thinking on, and compare accuracy, latency, output tokens, and cost.

Usage:
    export AI_BUILDER_TOKEN=sk_...        # Your AI Builder Space API key
    python run.py                          # Default: grok-4-fast, 16 sample tasks, 3 runs per mode
    python run.py --model gpt-5 --repeats 2
    python run.py --tasks my_tasks.jsonl   # Swap in your own tasks

Each line in the task file is one JSON object: {"id": "...", "question": "...", "answer": ["accepted answer 1", "answer 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')

# USD per 1M tokens (input, output). AI Builder Space rates as of 2026-10; update here if you change models or prices change.
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),
}

# The two modes differ only in reasoning_effort; model, tasks, and prompt stay identical.
# Both grok-4-fast and gpt-5 accept none (thinking off) as well as 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 fraction
    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):
    """Take the text after the last answer prefix; fall back to the last line if the model skipped the format."""
    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(',', ''))   # Pull out the number from replies like "130 degrees" or "$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,             # Thinking consumes output tokens; leave enough headroom or the reply may come back empty
        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='Number of full passes per mode')
    ap.add_argument('--workers', type=int, default=4, help='Number of concurrent requests')
    ap.add_argument('--out', default='results')
    args = ap.parse_args()

    token = os.getenv('AI_BUILDER_TOKEN')
    if not token:
        raise SystemExit('Please set the AI_BUILDER_TOKEN environment variable (your AI Builder Space API 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:                                   # Record failures too and count them as wrong
            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'Model {args.model}, {len(tasks)} tasks, {args.repeats} runs per mode, {len(jobs)} API calls total...')
    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## Summary ({args.model}, thinking off = reasoning_effort {MODES["off"]}, thinking on = {MODES["on"]})\n')
    print('| Mode | Accuracy (per run) | Median latency | Avg output tokens | Cost per full pass |')
    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 = 'Thinking off' if mode == 'off' else 'Thinking on'
        print(f'| {label} | {statistics.mean(per_rep):.1f}/{len(tasks)} ({", ".join(map(str, per_rep))}) '
              f'| {statistics.median(secs):.1f}s | {out_tok:.0f} | ${cost:.4f} |')

    print('\n## Correct runs per task (off → on)\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 = '  ← thinking helped' if on > off else ('  ← thinking hurt' 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'\nPer-item results (including the end of each model reply) saved to {path}')


if __name__ == '__main__':
    main()