"""5.1 Data for the temperature distribution chart: extract full-vocabulary step-1 logits for a prompt and apply softmax across temperatures.

Usage:
    python temperature_step1.py                       # Default prompt from logprobs.py (DEFAULT_PROMPTS[0])
    python temperature_step1.py --temps 0.5 1.0 1.5 --top 5

Loads the model the same way as logprobs.py (bf16, thinking mode disabled). Saves output to sample_output/temperature_step1_<device>.json:
raw logits for the top candidates, plus their probabilities and the combined tail probability ("rest") at each temperature.
"""
import argparse
import json
from datetime import datetime
from pathlib import Path

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

from logprobs import DEFAULT_PROMPTS, pick_device, show


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--model', default='Qwen/Qwen3.5-2B')
    ap.add_argument('--prompt', default=DEFAULT_PROMPTS[0])
    ap.add_argument('--temps', type=float, nargs='+', default=[0.5, 1.0, 1.5])
    ap.add_argument('--top', type=int, default=5, help='Number of top candidates to list individually')
    ap.add_argument('--out', default='sample_output')
    args = ap.parse_args()

    device = pick_device()
    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()
    msgs = [{'role': 'user', 'content': args.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)
    with torch.no_grad():
        logits = model(**inputs).logits[0, -1].float().cpu()          # Step 1: raw scores across the entire vocabulary

    order = torch.argsort(logits, descending=True)[:args.top]
    rec = dict(model=args.model, device=device, date=f'{datetime.now():%Y-%m}', prompt=args.prompt, vocab=int(logits.shape[0]),
               top=[dict(token=show(tok, [int(i)]), id=int(i), logit=float(logits[i])) for i in order], temps={})
    for t in args.temps:
        p = (logits / t).softmax(-1)
        rec['temps'][str(t)] = dict(probs=[float(p[i]) for i in order], rest=float(1 - p[order].sum()))
    print(f'Model {args.model}, device {device}, prompt: {args.prompt}')
    print('Candidate | logit | ' + ' | '.join(f'Temperature {t}' for t in args.temps))
    for k, c in enumerate(rec['top']):
        print(f'{c["token"]} | {c["logit"]:.2f} | ' + ' | '.join(f'{rec["temps"][str(t)]["probs"][k]:.3f}' for t in args.temps))
    print('All remaining | | ' + ' | '.join(f'{rec["temps"][str(t)]["rest"]:.3f}' for t in args.temps))
    Path(args.out).mkdir(exist_ok=True)
    path = Path(args.out) / f'temperature_step1_{device}.json'
    path.write_text(json.dumps(rec, ensure_ascii=False, indent=1), encoding='utf-8')
    print(f'Saved to {path}')


if __name__ == '__main__':
    main()