"""
generate.py -- inference / generation harness for Quazimoto-LM.

Loads a checkpoint saved by train.py (save_ckpt writes
    {"model", "family_config", "vocab_size", "step"})
rebuilds the exact model from the embedded family_config, restores weights, and
samples completions with the bundled SpikeWhale tokenizer.

This is tuned to the Quazimoto architecture's quirks:
  * The model has NO KV cache (attention is training-only / full recompute), so
    each new token re-runs the whole context. We respect cfg.block_size by
    cropping the context window, exactly like model.generate().
  * We reuse the model's NaN/inf-sanitised sampling path (the oscillator stack
    can push logits to +/-inf over long rollouts), and add nucleus (top-p),
    top-k, and repetition-penalty controls on top.
  * Generation stops early on <eos> / <|im_end|> / <|endoftext|> when present.

Examples
--------
  # plain completion
  python generate.py --ckpt chkpt/quazimoto.pt --prompt "the quazimoto oscillator"

  # chat-style turn (wraps prompt in ChatML and stops on <|im_end|>)
  python generate.py --ckpt chkpt/quazimoto.pt --chat --prompt "Hello, who are you?"

  # interactive REPL
  python generate.py --ckpt chkpt/quazimoto.pt --interactive
"""

import argparse
import os
import sys

import torch
import torch.nn.functional as F

from model import QuazimotoLM, QuazimotoConfig

PKG_DIR = os.path.dirname(os.path.abspath(__file__))


# --------------------------------------------------------------------------- #
# loading
# --------------------------------------------------------------------------- #
def load_tokenizer(tok_dir):
    sys.path.insert(0, tok_dir)
    from spike_tokenizer import SpikeTokenizer
    return SpikeTokenizer(vocab_file=os.path.join(tok_dir, "tokenizer.json"))


def load_model(ckpt_path, device):
    ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
    fam = ckpt.get("family_config")
    if fam is None:
        raise ValueError(f"{ckpt_path} has no family_config; can't rebuild the model.")
    fam = dict(fam)
    state = ckpt["model"]

    # Morphable-hybrid detection. Older hybrid checkpoints saved family_config WITHOUT
    # use_morphable (the attention weights live under `.attn.full.*` and there are linear
    # twins under `.attn.lin.*`). If we rebuilt a plain model those weights would all be
    # "missing" -> random attention -> gibberish. Force the morphable/triattn structure
    # so the state dict maps correctly, then apply the discretized modes + compressor.
    is_hybrid = ("modes" in ckpt) or any(".attn.lin." in k for k in state)
    if is_hybrid:
        fam["use_morphable"] = True
        fam.setdefault("morph_k_full", (ckpt.get("morph") or {}).get("k_full", 4))
        tri = ckpt.get("triattn") or {}
        if ckpt.get("triattn_stats") is not None or tri:
            fam["use_triattention"] = True
            fam.setdefault("triattn_budget", tri.get("budget", 2048))
            fam.setdefault("triattn_window", tri.get("window", 128))

    cfg = QuazimotoConfig(**fam)
    model = QuazimotoLM(cfg)
    missing, unexpected = model.load_state_dict(state, strict=False)
    if missing:
        print(f"  [warn] missing keys: {len(missing)} (e.g. {missing[:3]})")
    if unexpected:
        print(f"  [warn] unexpected keys: {len(unexpected)} (e.g. {unexpected[:3]})")

    if is_hybrid:
        modes = ckpt.get("modes")
        if modes is not None:
            for layer, mode in zip(model.layers, modes):
                layer.attn.mode = mode
        if ckpt.get("triattn_stats") is not None and model.triattn is not None:
            model.triattn.load_state_dict(ckpt["triattn_stats"])
        model._is_hybrid = True
        print(f"  [hybrid] modes={ckpt.get('modes')} | kept full={ckpt.get('kept_full_layers')}")

    model.to(device).eval()
    step = ckpt.get("step", "?")
    print(f"loaded {ckpt_path} | step {step} | vocab {cfg.vocab_size} | "
          f"block_size {cfg.block_size}")
    return model, cfg


# --------------------------------------------------------------------------- #
# sampling -- tuned for the oscillator stack (no KV cache, sanitised logits)
# --------------------------------------------------------------------------- #
def _filter_logits(logits, top_k, top_p):
    """Apply top-k then nucleus (top-p) filtering in place; returns logits."""
    if top_k:
        k = min(top_k, logits.size(-1))
        kth = torch.topk(logits, k).values[..., -1, None]
        logits = logits.masked_fill(logits < kth, float("-inf"))
    if top_p and 0.0 < top_p < 1.0:
        sorted_logits, sorted_idx = torch.sort(logits, descending=True, dim=-1)
        cum = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
        remove = cum > top_p
        remove[..., 1:] = remove[..., :-1].clone()  # keep the first over-threshold token
        remove[..., 0] = False
        scatter_remove = remove.scatter(-1, sorted_idx, remove)
        logits = logits.masked_fill(scatter_remove, float("-inf"))
    return logits


@torch.no_grad()
def generate(model, cfg, idx, n_new, temperature=0.9, top_k=40, top_p=0.95,
             repetition_penalty=1.1, stop_ids=None, device="cpu", use_cache=True):
    """Autoregressive sampling. Mirrors model.generate()'s NaN/inf sanitisation
    (long oscillator rollouts can emit +/-inf logits whose softmax -> NaN ->
    multinomial device-side assert) and adds top-p + repetition penalty.

    Uses the attention KV cache by default: the prompt is prefilled once, then
    each new token is a single-position forward. Once the cache reaches
    block_size we re-prime with a windowed recompute (absolute RoPE positions
    must stay within max_position_embeddings)."""
    model.eval()
    stop_ids = set(stop_ids or [])
    past = None
    for _ in range(n_new):
        if use_cache and past is not None and past[0][0].size(2) < cfg.block_size:
            logits, _, _, past = model(idx[:, -1:], past_key_values=past, use_cache=True)
        elif use_cache:
            logits, _, _, past = model(idx[:, -cfg.block_size:], use_cache=True)
        else:
            logits, _, _ = model(idx[:, -cfg.block_size:])
        lg = torch.nan_to_num(logits[:, -1, :].float(),
                              nan=0.0, posinf=1e4, neginf=-1e4)

        # repetition penalty over tokens already in the context
        if repetition_penalty and repetition_penalty != 1.0:
            for b in range(lg.size(0)):
                seen = torch.unique(idx[b])
                vals = lg[b, seen]
                lg[b, seen] = torch.where(vals > 0, vals / repetition_penalty,
                                          vals * repetition_penalty)

        lg = lg / max(temperature, 1e-6)
        lg = _filter_logits(lg, top_k, top_p)
        probs = torch.softmax(lg, dim=-1)

        if not torch.isfinite(probs).all() or float(probs.sum()) <= 0.0:
            nxt = torch.argmax(torch.nan_to_num(lg, neginf=-1e4), dim=-1, keepdim=True)
        else:
            nxt = torch.multinomial(probs, 1)

        idx = torch.cat([idx, nxt], dim=1)
        if stop_ids and int(nxt[0]) in stop_ids and idx.size(0) == 1:
            break
    return idx


@torch.no_grad()
def generate_anti_drift(model, cfg, idx, n_new, *, temperature=0.7, min_p=0.05,
                        top_k=0, no_repeat_ngram=3, cfg_gamma=1.35,
                        repetition_penalty=1.1, stop_ids=None, device="cpu"):
    """Anti-drift decoding for small models (tuned in decode_lab.py).

    Combines three levers that each target a cause of coherence drift:
      * CFG / context guidance (cfg_gamma): run a second UNCONDITIONAL stream that
        withholds the prompt and push logits away from it,
            logits = logits_uncond + gamma*(logits_cond - logits_uncond),
        so the model commits harder to what the prompt implies. Keep gamma <=1.5;
        above ~2 it over-amplifies and mangles word surface form.
      * min_p: drop tokens below min_p * p(top) -- a dynamic tail cut that beats
        top-p on the flat next-token distributions small models produce.
      * no_repeat_ngram + repetition_penalty: kill the loop/degeneration mode.

    Batch-1. Uses a KV cache for each stream; re-primes on block_size overflow."""
    assert idx.size(0) == 1, "anti-drift decoder is batch-1"
    model.eval()
    stop_ids = set(stop_ids or [])
    use_cfg = abs(cfg_gamma - 1.0) > 1e-6
    B = cfg.block_size

    def step(seq, past):
        if past is not None and past[0][0].size(2) < B:
            lo, _, _, past = model(seq[:, -1:], past_key_values=past, use_cache=True)
        else:
            lo, _, _, past = model(seq[:, -B:], use_cache=True)
        return lo[:, -1, :].float(), past

    past_c = past_u = None
    useq = idx[:, -1:].clone() if use_cfg else None   # uncond stream withholds prompt

    seq = idx
    for _ in range(n_new):
        logits_c, past_c = step(seq, past_c)           # 1st iter prefills the prompt
        if use_cfg:
            logits_u, past_u = step(useq, past_u)
            lg = logits_u + cfg_gamma * (logits_c - logits_u)
        else:
            lg = logits_c
        lg = torch.nan_to_num(lg, nan=0.0, posinf=1e4, neginf=-1e4)

        if repetition_penalty and repetition_penalty != 1.0:
            seen = torch.unique(seq[0])
            vals = lg[0, seen]
            lg[0, seen] = torch.where(vals > 0, vals / repetition_penalty,
                                      vals * repetition_penalty)
        if no_repeat_ngram and seq.size(1) >= no_repeat_ngram - 1:
            n = no_repeat_ngram
            ids_list = seq[0].tolist()
            prefix = tuple(ids_list[-(n - 1):]) if n > 1 else ()
            for i in range(len(ids_list) - n + 1):
                if tuple(ids_list[i:i + n - 1]) == prefix:
                    lg[0, ids_list[i + n - 1]] = -float("inf")

        lg = lg / max(temperature, 1e-6)
        if min_p and min_p > 0:
            p = torch.softmax(lg, dim=-1)
            lg = lg.masked_fill(p < min_p * p.max(dim=-1, keepdim=True).values, -float("inf"))
        lg = _filter_logits(lg, top_k, 1.0)

        probs = torch.softmax(lg, dim=-1)
        if not torch.isfinite(probs).all() or float(probs.sum()) <= 0.0:
            nxt = torch.argmax(torch.nan_to_num(lg, neginf=-1e4), dim=-1, keepdim=True)
        else:
            nxt = torch.multinomial(probs, 1)
        seq = torch.cat([seq, nxt], dim=1)
        if use_cfg:
            useq = torch.cat([useq, nxt], dim=1)
        if stop_ids and int(nxt[0]) in stop_ids:
            break
    return seq


# --------------------------------------------------------------------------- #
# self-speculative decoding (DeepSpec draft+verify, drafted by the MTP heads)
# --------------------------------------------------------------------------- #
@torch.no_grad()
def generate_speculative(model, cfg, idx, n_new, stop_ids=None, verbose=False, trace=None):
    """Greedy self-speculative decoding (batch size 1).

    DeepSpec-style draft+verify, but the model drafts on ITSELF: the `mtp_layers`
    MTP heads propose the next `mtp_layers` future tokens from one hidden state,
    and the main head verifies them in a single parallel forward, accepting the
    longest correct prefix. Output is BIT-IDENTICAL to greedy generate() (top_k=1)
    -- speculation only changes the number of forwards, never the tokens.

    Per cycle (one forward over committed seq + the pending drafts):
      base = index of the last committed token; main_logits[base+j] is the
      verifier's genuine token AFTER draft j. We accept draft j while it matches,
      commit one correction token at the first mismatch, then read the next round's
      drafts from the MTP heads at the deepest still-valid position (base+a)."""
    assert idx.size(0) == 1, "speculative decoder is batch-1"
    model.eval()
    if model.mtp_heads is None:
        raise ValueError("speculative decoding needs MTP heads (use_mtp=True).")
    stop_ids = set(stop_ids or [])
    K = len(model.mtp_heads)
    start_len = idx.size(1)
    forwards = 0

    def crop(seq):                                   # honor block_size context window
        return seq[:, -cfg.block_size:]

    # bootstrap: one forward over the prompt -> genuine next token + K drafts
    main_logits, mtp = model.forward_drafts(crop(idx)); forwards += 1
    g = main_logits[:, -1].argmax(-1, keepdim=True)
    idx = torch.cat([idx, g], dim=1)
    drafts = [mtp[k][:, -1].argmax(-1, keepdim=True) for k in range(K)]  # guesses after g

    while idx.size(1) - start_len < n_new:
        if stop_ids and int(idx[0, -1]) in stop_ids:
            break
        seq = torch.cat([idx] + drafts, dim=1)        # committed + K speculative tokens
        base = idx.size(1) - 1                         # position of the last committed token
        cseq = crop(seq)
        off = cseq.size(1) - seq.size(1)              # crop shift (<=0) to realign `base`
        main_logits, mtp = model.forward_drafts(cseq); forwards += 1
        b = base + off

        # verify: accept draft[j] while it matches the verifier's token after it
        a = 0
        for j in range(K):
            v = main_logits[:, b + j].argmax(-1, keepdim=True)   # genuine token after draft j-1 / g
            if int(v) == int(drafts[j]):
                a += 1
            else:
                break
        accepted = drafts[:a]
        correction = main_logits[:, b + a].argmax(-1, keepdim=True)  # genuine token at first miss
        new_toks = accepted + [correction]
        idx = torch.cat([idx] + new_toks, dim=1)
        if trace is not None:                          # per-cycle record for the visualizer
            trace.append({"drafts": [int(d) for d in drafts], "accepted": a,
                          "correction": int(correction)})

        # next drafts from the deepest still-valid hidden (position b+a holds the last
        # accepted token / g, whose MTP heads predict the tokens AFTER `correction`).
        drafts = [mtp[k][:, b + a].argmax(-1, keepdim=True) for k in range(K)]

        if stop_ids and any(int(t) in stop_ids for t in new_toks):
            break

    idx = idx[:, :start_len + n_new]
    if verbose:
        produced = idx.size(1) - start_len
        print(f"  [spec] {produced} tokens in {forwards} forwards "
              f"({produced / max(forwards,1):.2f} tok/forward; "
              f"vs 1.00 for plain greedy)")
    return idx


# --------------------------------------------------------------------------- #
# prompt helpers
# --------------------------------------------------------------------------- #
def build_prompt(tok, text, chat, system):
    """Return token ids for the prompt. In chat mode, wrap in ChatML so the
    model is cued to produce an assistant turn (only useful if it was trained
    on that framing -- harmless otherwise)."""
    if chat:
        parts = []
        if system:
            parts.append(f"<|im_start|><|system|>\n{system}<|im_end|>\n")
        parts.append(f"<|im_start|><|user|>\n{text}<|im_end|>\n")
        parts.append("<|im_start|><|assistant|>\n")
        text = "".join(parts)
    ids = tok.encode(text, add_special_tokens=False)
    return ids, text


def resolve_stop_ids(tok, chat):
    """Collect ids of end-of-turn / end-of-text markers that exist in the vocab."""
    names = ["<eos>", "<|endoftext|>"]
    if chat:
        names.insert(0, "<|im_end|>")
    vocab = tok.get_vocab()
    return [vocab[n] for n in names if n in vocab]


# --------------------------------------------------------------------------- #
# main
# --------------------------------------------------------------------------- #
def run_once(model, cfg, tok, args, device, prompt_text):
    stop_ids = resolve_stop_ids(tok, args.chat)
    ids, shown = build_prompt(tok, prompt_text, args.chat, args.system)
    x = torch.tensor([ids], dtype=torch.long, device=device)
    if getattr(model, "_is_hybrid", False):
        # morphable hybrid: linear layers carry O(1) state, full layers a KV cache the
        # TriAttention compressor prunes. Must use the hybrid decode path.
        out = model.generate_hybrid(x, args.max_new_tokens,
                                    temperature=args.temperature, top_k=args.top_k)
    elif args.speculative:
        # greedy self-speculative decoding (MTP heads draft, main head verifies);
        # output equals greedy decoding, only faster. Ignores sampling knobs.
        out = generate_speculative(model, cfg, x, args.max_new_tokens,
                                   stop_ids=stop_ids, verbose=args.spec_stats)
    elif getattr(args, "anti_drift", False):
        out = generate_anti_drift(model, cfg, x, args.max_new_tokens,
                                  temperature=args.temperature, min_p=args.min_p,
                                  top_k=args.top_k or 0, no_repeat_ngram=args.no_repeat_ngram,
                                  cfg_gamma=args.cfg_gamma,
                                  repetition_penalty=args.repetition_penalty,
                                  stop_ids=stop_ids, device=device)
    else:
        out = generate(model, cfg, x, args.max_new_tokens,
                       temperature=args.temperature, top_k=args.top_k,
                       top_p=args.top_p, repetition_penalty=args.repetition_penalty,
                       stop_ids=stop_ids, device=device, use_cache=not args.no_cache)
    gen_ids = out[0, len(ids):].tolist()
    completion = tok.decode(gen_ids, skip_special_tokens=not args.show_special)
    if args.echo_prompt:
        print(shown, end="")
    print(completion)
    return completion


def main():
    p = argparse.ArgumentParser(description="Quazimoto-LM generation harness")
    p.add_argument("--ckpt", default=os.path.join(PKG_DIR, "chkpt", "quazimoto.pt"),
                   help="checkpoint .pt saved by train.py")
    p.add_argument("--tok_dir", default=PKG_DIR, help="dir holding tokenizer.json")
    p.add_argument("--prompt", default="the quazimoto oscillator")
    p.add_argument("--max_new_tokens", type=int, default=200)
    p.add_argument("--temperature", type=float, default=0.9)
    p.add_argument("--top_k", type=int, default=40, help="0 to disable")
    p.add_argument("--top_p", type=float, default=0.95, help="1.0 to disable")
    p.add_argument("--repetition_penalty", type=float, default=1.1, help="1.0 to disable")
    p.add_argument("--seed", type=int, default=None)
    p.add_argument("--chat", action="store_true", help="wrap prompt in ChatML framing")
    p.add_argument("--system", default="", help="system prompt (chat mode only)")
    p.add_argument("--echo_prompt", action="store_true", help="print the prompt before output")
    p.add_argument("--show_special", action="store_true", help="don't strip special tokens")
    p.add_argument("--interactive", action="store_true", help="REPL: read prompts from stdin")
    p.add_argument("--no_cache", action="store_true",
                   help="disable the KV cache (full recompute each step)")
    p.add_argument("--anti-drift", dest="anti_drift", action="store_true",
                   help="anti-drift decoding for small models: CFG context guidance + "
                        "min_p + no-repeat-ngram (tuned in decode_lab.py)")
    p.add_argument("--cfg_gamma", type=float, default=1.35,
                   help="anti-drift CFG strength (1=off; keep <=1.5, >2 mangles words)")
    p.add_argument("--min_p", type=float, default=0.05, help="anti-drift min_p tail cut")
    p.add_argument("--no_repeat_ngram", type=int, default=3, help="anti-drift n-gram block (0=off)")
    p.add_argument("--speculative", action="store_true",
                   help="greedy self-speculative decoding via MTP draft heads (DeepSpec-style)")
    p.add_argument("--spec_stats", action="store_true",
                   help="print tokens/forward acceptance stats for speculative decoding")
    p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    args = p.parse_args()

    if args.top_k == 0:
        args.top_k = None
    if args.seed is not None:
        torch.manual_seed(args.seed)

    device = args.device
    tok = load_tokenizer(args.tok_dir)
    model, cfg = load_model(args.ckpt, device)

    if cfg.vocab_size != tok.vocab_size:
        print(f"  [warn] model vocab {cfg.vocab_size} != tokenizer vocab "
              f"{tok.vocab_size}; decode may be misaligned.")

    if args.interactive:
        print("interactive mode -- type a prompt, Ctrl-C / empty line + EOF to quit.\n")
        try:
            while True:
                try:
                    line = input(">>> ").strip()
                except EOFError:
                    break
                if not line:
                    continue
                run_once(model, cfg, tok, args, device, line)
                print()
        except KeyboardInterrupt:
            print("\nbye.")
    else:
        run_once(model, cfg, tok, args, device, args.prompt)


if __name__ == "__main__":
    main()