"""Train-aligned generation for Fractus CTE. No frequency ban / scramble."""
from __future__ import annotations
from typing import List, Optional
import torch

def _reset(engine):
    engine.eval()
    if hasattr(engine, "reset_thought"):
        engine.reset_thought(1)
    for blk in getattr(engine, "blocks", []):
        if hasattr(blk, "attn_S"):
            blk.attn_S.zero_()
        if hasattr(blk, "attn_z"):
            blk.attn_z.zero_()

def _pick(logits: torch.Tensor, prev: Optional[int], temperature: float, top_k: int) -> int:
    l = logits.float().reshape(-1).clone()
    if prev is not None and 0 <= prev < l.numel():
        l[prev] = -1e9
    if temperature <= 1e-5:
        return int(l.argmax().item())
    l = l / max(temperature, 1e-5)
    k = min(max(1, top_k), l.numel())
    topv, topi = torch.topk(l, k)
    return int(topi[torch.multinomial(torch.softmax(topv, -1), 1)].item())

@torch.no_grad()
def generate_window(
    engine,
    tokenizer,
    prompt: str,
    max_new: int = 40,
    temperature: float = 0.0,
    top_k: int = 40,
    window: int = 128,
    force_second: bool = False,
    carry: bool = False,
) -> tuple[str, List[int]]:
    ids = tokenizer.encode(prompt)[:window] or [0]
    out: List[int] = []
    prev = ids[-1]
    if carry:
        _reset(engine)
        logits = engine.tick_chunk(torch.tensor([ids], dtype=torch.long))
        cur = logits[0, -1]
        for i in range(max_new):
            if force_second and i == 0:
                l = cur.float().clone()
                l[int(l.argmax())] = -1e9
                nxt = _pick(l, prev=prev, temperature=temperature, top_k=top_k)
            else:
                nxt = _pick(cur, prev=prev, temperature=temperature, top_k=top_k)
            out.append(nxt)
            ids.append(nxt)
            logits = engine.tick_chunk(torch.tensor([[nxt]], dtype=torch.long))
            cur = logits[0, -1]
            prev = nxt
    else:
        for i in range(max_new):
            _reset(engine)
            ctx = ids[-window:]
            logits = engine.tick_chunk(torch.tensor([ctx], dtype=torch.long))
            cur = logits[0, -1]
            if force_second and i == 0:
                l = cur.float().clone()
                l[int(l.argmax())] = -1e9
                nxt = _pick(l, prev=prev, temperature=temperature, top_k=top_k)
            else:
                nxt = _pick(cur, prev=prev, temperature=temperature, top_k=top_k)
            out.append(nxt)
            ids.append(nxt)
            prev = nxt
    return tokenizer.decode(out), out

@torch.no_grad()
def generate_greedy_prefix(engine, tokenizer, prompt: str, max_new: int = 40, window: int = 128):
    return generate_window(engine, tokenizer, prompt, max_new=max_new, temperature=0.0, top_k=1, window=window, carry=False)

@torch.no_grad()
def generate_greedy_ids(engine, tokenizer, prompt: str, max_new: int = 40, window: int = 128):
    return generate_greedy_prefix(engine, tokenizer, prompt, max_new=max_new, window=window)

@torch.no_grad()
def generate_chunk(engine, tokenizer, prompt: str, max_new: int = 40, **kw):
    return generate_window(engine, tokenizer, prompt, max_new=max_new, window=128, **kw)

@torch.no_grad()
def unique40_probe(engine, max_new=40, mode="prefix", prompts=None):
    return {"mode": mode, "max_new": max_new}

# space-sync 2026-08-29b force_second carry