"""generate.py — запуск обученной reasoning-модели по чекпойнту.

Даёшь задачу по коду/математике -> модель рассуждает в <think>...</think> и
выдаёт ответ. Декодинг: полный пересчёт растущей последовательности (просто и
корректно для hybrid: и GLA-, и attn-слои работают через обычный forward).

Запуск:
  python generate.py --ckpt out_reason/ckpt_last.pt --prompt "Напиши функцию..."
  python generate.py --ckpt out_reason/ckpt_last.pt           # интерактивно
  python generate.py --ckpt out_reason/ckpt_last.pt --no_think  # без затравки <think>
"""

import argparse, os
import torch
import torch.nn.functional as F

from codetrain.model import CodeLM, ModelConfig

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
THINK, ENDTHINK = "<think>", "</think>"


def load_model(ckpt_path):
    ck = torch.load(ckpt_path, map_location=DEVICE)
    cfg = ModelConfig(**ck["cfg"])
    cfg.grad_checkpoint = False                 # инференс: без чекпойнтинга
    model = CodeLM(cfg).to(DEVICE).to(torch.bfloat16)
    model.load_state_dict(ck["model"])
    model.eval()
    step = ck.get("step", "?")
    print(f"Загружен чекпойнт: шаг {step} | {model.num_params()/1e6:.0f}M "
          f"| ctx {cfg.block_size} | mixer {cfg.mixer}")
    return model, cfg


def build_tokenizer(name="bigcode/starcoder2-15b"):
    from transformers import AutoTokenizer
    tok = AutoTokenizer.from_pretrained(name)
    if tok.eos_token_id is None:
        tok.add_special_tokens({"eos_token": "<|endoftext|>"})
    add = [t for t in (THINK, ENDTHINK) if t not in tok.get_vocab()]
    if add:
        tok.add_special_tokens({"additional_special_tokens": add})
    return tok


def _sample(logits, temperature, top_p):
    """Один токен из распределения. temperature=0 -> жадно (argmax)."""
    if temperature <= 0:
        return int(logits.argmax(-1))
    logits = logits / temperature
    probs = F.softmax(logits, dim=-1)
    if 0 < top_p < 1:                            # nucleus: оставляем ядро массы top_p
        sp, si = torch.sort(probs, descending=True)
        csum = torch.cumsum(sp, dim=-1)
        keep = csum - sp <= top_p                # держим, пока кумулятив не превысил
        sp = sp * keep
        sp = sp / sp.sum()
        nxt = si[torch.multinomial(sp, 1)]
        return int(nxt)
    return int(torch.multinomial(probs, 1))


@torch.no_grad()
def generate(model, cfg, tok, prompt, max_new=512, temperature=0.7, top_p=0.95,
             seed_think=True, stream=True):
    """Сгенерировать продолжение. seed_think=True добавляет '<think>\\n' к затравке,
    подталкивая модель сначала рассуждать. Останов по EOS. Печатает по токенам."""
    text = prompt.rstrip() + "\n"
    if seed_think:
        text += THINK + "\n"
    ids = tok.encode(text)
    if stream:
        print("\n--- генерация ---")
        print(text, end="", flush=True)
    ids = torch.tensor([ids], device=DEVICE)
    eos = tok.eos_token_id
    out_ids = []
    for _ in range(max_new):
        ctx = ids[:, -cfg.block_size:]           # не выходим за обученный контекст
        logits, _ = model(ctx)                   # (1, 1, vocab) — последний шаг
        nxt = _sample(logits[0, -1], temperature, top_p)
        if nxt == eos:
            break
        out_ids.append(nxt)
        ids = torch.cat([ids, torch.tensor([[nxt]], device=DEVICE)], dim=1)
        if stream:
            print(tok.decode([nxt]), end="", flush=True)
    if stream:
        print("\n--- конец ---")
    return tok.decode(out_ids)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", required=True)
    ap.add_argument("--tokenizer", default="bigcode/starcoder2-15b")
    ap.add_argument("--prompt", default=None, help="одноразовый запрос (иначе интерактив)")
    ap.add_argument("--max_new", type=int, default=512)
    ap.add_argument("--temperature", type=float, default=0.7)
    ap.add_argument("--top_p", type=float, default=0.95)
    ap.add_argument("--no_think", action="store_true",
                    help="не добавлять <think> в затравку (модель сама решит)")
    args = ap.parse_args()

    if not os.path.exists(args.ckpt):
        print(f"Чекпойнт не найден: {args.ckpt}"); return
    tok = build_tokenizer(args.tokenizer)
    model, cfg = load_model(args.ckpt)

    def run(p):
        generate(model, cfg, tok, p, max_new=args.max_new,
                 temperature=args.temperature, top_p=args.top_p,
                 seed_think=not args.no_think, stream=True)

    if args.prompt is not None:
        run(args.prompt)
    else:
        print("Интерактив. Пустая строка / Ctrl-C — выход.")
        while True:
            try:
                p = input("\n>>> запрос: ").strip()
            except (EOFError, KeyboardInterrupt):
                print(); break
            if not p:
                break
            run(p)


if __name__ == "__main__":
    main()