"""5.1 How each token is chosen: inspect per-step candidate tokens and probabilities, then sample across temperatures.

Usage:
    python logprobs.py                                   # Two default prompts, default model Qwen/Qwen3.5-2B
    python logprobs.py --prompt "Your own question"
    python logprobs.py --model Qwen/Qwen3.5-0.8B         # Smaller model for machines with less memory
    python logprobs.py --temps 0.3 1.0 1.5 --samples 5

Automatically uses CUDA on NVIDIA GPUs, MPS on Apple Silicon Macs, or CPU otherwise (slow, but works).
"""
import argparse
import json
from datetime import datetime
from pathlib import Path

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

DEFAULT_PROMPTS = [
    '法国的首都是哪座城市？用一句话回答。',
    '推荐一家上海的独立咖啡馆，给出店名和具体地址，一句话回答。',
]


def pick_device():
    if torch.cuda.is_available():
        return 'cuda'
    if torch.backends.mps.is_available():
        return 'mps'
    return 'cpu'


def show(tok, ids):
    """Format a token for display: show spaces as · and newlines as ⏎."""
    return tok.decode(ids).replace(' ', '·').replace('\n', '⏎') or '∅'


def generate(model, tok, prompt, temperature, max_new_tokens, n=1, top_p=1.0):
    """Generate n replies. Return token IDs for each reply and per-step probabilities from raw logits (unscaled model distribution)."""
    msgs = [{'role': 'user', 'content': prompt}]
    text = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True, enable_thinking=False)
    inputs = tok(text, return_tensors='pt').to(model.device)
    kw = dict(do_sample=False) if temperature == 0 else dict(do_sample=True, temperature=temperature, top_p=top_p, top_k=0)
    out = model.generate(**inputs, max_new_tokens=max_new_tokens, num_return_sequences=n, output_logits=True,
                         return_dict_in_generate=True, pad_token_id=tok.eos_token_id, **kw)
    seqs = out.sequences[:, inputs['input_ids'].shape[1]:]
    probs = torch.stack(out.logits, dim=1).float().softmax(-1)          # [n, steps, vocab]
    return seqs, probs


def step_table(tok, seq, probs, low, top_k=5, top_p=0.9):
    """Step-by-step table for greedy decoding: chosen token, its probability, runner-up candidates, and candidates needed to reach top_p."""
    rows = []
    for i, tid in enumerate(seq.tolist()):
        if tid == tok.eos_token_id:
            break
        p = probs[i]
        top = torch.topk(p, top_k)
        nucleus = int((torch.cumsum(torch.sort(p, descending=True).values, 0) < top_p).sum()) + 1
        rows.append(dict(step=i + 1, token=show(tok, [tid]), prob=float(p[tid]), nucleus=nucleus,
                         alternatives=[(show(tok, [int(j)]), float(q)) for q, j in zip(top.values, top.indices) if int(j) != tid][:top_k - 1]))
    print('| Step | Chosen token | Probability | Runner-up candidates (probability) | Candidates needed for 90% mass |')
    print('|---|---|---|---|---|')
    for r in rows:
        flag = ' ⚠' if r['prob'] < low else ''
        alts = ', '.join(f'{t} {q:.2f}' for t, q in r['alternatives'])
        print(f'| {r["step"]} | {r["token"]} | {r["prob"]:.2f}{flag} | {alts} | {r["nucleus"]} |')
    return rows


def marked(tok, seq, probs, low):
    """Mark tokens with probability below low in the reply using 【token|prob】."""
    parts = []
    for i, tid in enumerate(seq.tolist()):
        if tid == tok.eos_token_id:
            break
        s = tok.decode([tid])
        p = float(probs[i][tid])
        parts.append(f'【{s}|{p:.2f}】' if p < low else s)
    return ''.join(parts).replace('\n', ' ')


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--model', default='Qwen/Qwen3.5-2B')
    ap.add_argument('--prompt', action='append', help='Can be passed multiple times; defaults to two built-in prompts')
    ap.add_argument('--temps', type=float, nargs='+', default=[0.3, 0.7, 1.0, 1.5])
    ap.add_argument('--samples', type=int, default=5, help='Number of samples per temperature')
    ap.add_argument('--top-p', type=float, default=0.9, help='Run one extra round at --top-p-temp with this top-p for comparison')
    ap.add_argument('--top-p-temp', type=float, default=1.0)
    ap.add_argument('--max-new-tokens', type=int, default=48)
    ap.add_argument('--low', type=float, default=0.3, help='Highlight tokens whose probability falls below this threshold')
    ap.add_argument('--seed', type=int, default=0)
    ap.add_argument('--out', default='results')
    args = ap.parse_args()

    device = pick_device()
    torch.manual_seed(args.seed)
    tok = AutoTokenizer.from_pretrained(args.model)
    model = AutoModelForCausalLM.from_pretrained(args.model, dtype=torch.float32 if device == 'cpu' else torch.bfloat16).to(device)
    model.eval()
    print(f'Model {args.model}, device {device}, {datetime.now():%Y-%m-%d %H:%M}.')
    print(f'Probabilities always show the raw model distribution (temperature 1); tokens below {args.low} are flagged with ⚠ or shown as 【token|prob】.')
    print('No top-k cutoff is applied during sampling; except for the final group, top-p is 1 (only temperature varies).')

    record = dict(model=args.model, device=device, prompts=[])
    for prompt in args.prompt or DEFAULT_PROMPTS:
        print(f'\n# Prompt: {prompt}\n')
        print('## Temperature 0 (pick the highest-probability token at each step)\n')
        seqs, probs = generate(model, tok, prompt, 0, args.max_new_tokens)
        rows = step_table(tok, seqs[0], probs[0], args.low)
        print(f'\nReply: {marked(tok, seqs[0], probs[0], args.low)}')
        rec = dict(prompt=prompt, greedy=rows, samples={})
        settings = [(t, 1.0) for t in args.temps] + ([(args.top_p_temp, args.top_p)] if args.top_p < 1 else [])
        for t, top_p in settings:
            label = f'温度 {t}' + (f' + top-p {top_p}' if top_p < 1 else '')
            print(f'\n## Temperature {t}' + (f' + top-p {top_p}' if top_p < 1 else '') + f', {args.samples} samples\n')
            seqs, probs = generate(model, tok, prompt, t, args.max_new_tokens, n=args.samples, top_p=top_p)
            texts = []
            for k in range(args.samples):
                m = marked(tok, seqs[k], probs[k], args.low)
                n_low = m.count('【')
                print(f'{k + 1}. {m} ({n_low} low-probability tokens)')
                texts.append(dict(text=tok.decode(seqs[k], skip_special_tokens=True), marked=m, n_low=n_low))
            print(f'\n{len({x["text"] for x in texts})} distinct replies across {args.samples} runs.')
            rec['samples'][label] = texts
        record['prompts'].append(rec)

    Path(args.out).mkdir(exist_ok=True)
    path = Path(args.out) / f'logprobs_{args.model.split("/")[-1]}_{device}_{datetime.now():%Y%m%d_%H%M}.json'
    path.write_text(json.dumps(record, ensure_ascii=False, indent=1), encoding='utf-8')
    print(f'\nFull results saved to {path}')


if __name__ == '__main__':
    main()