import os
import re
import warnings
import shutil
import sys
import threading
from pathlib import Path
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
import inspect
from dataclasses import dataclass
from typing import Optional, Any, List

# Global event: set this to interrupt generation mid-stream.
# The generation loop polls it every token. Cleared before each new generation.
STOP_GENERATION = threading.Event()

# Special Tokens
BOT_TOKEN = "<bot>"
EOT_TOKEN = "<eot>"
STEP_TOKEN = "<step>"

# Kept deliberately short. This model's instruction-following is weak, and a
# long policy preamble measurably crowds out the actual question. Two clauses
# is about the most it reliably holds.
SYSTEM_PROMPT = (
    # The first sentence is VERBATIM the system prompt the identity examples
    # were trained under. Replacing it wholesale moved those rows
    # off-distribution and the model started emitting a visible <think> block
    # instead of answering "who are you". Extra instructions go AFTER it.
    "You are Hyper, a helpful and cooperative assistant. Follow "
    "instructions directly and refer to yourself as Hyper. "
    "You were created by Cymela. "
    "Refuse to help with illegal activities."
)

# ---------------------------------------------------------------------------
# Pretty printing: markdown + LaTeX -> ANSI + unicode
#
# The model emits markdown (**bold**, ###, `code`) and LaTeX ($$..$$,
# \text{}, p_{N+1}, \times) because that is what its training corpus looks
# like.  Raw, that is unreadable in a terminal.  This renders it in place.
# Toggle at runtime with /render on|off.
# ---------------------------------------------------------------------------
BOLD, DIM, ITAL, CYAN, YELL, GREEN, RESET = (
    "\033[1m", "\033[2m", "\033[3m", "\033[36m", "\033[33m", "\033[32m", "\033[0m")


def enable_ansi() -> bool:
    """Turn on VT processing on Windows consoles; no-op elsewhere."""
    if os.name != "nt":
        return sys.stdout.isatty()
    try:
        import ctypes
        kernel32 = ctypes.windll.kernel32
        # -11 = STD_OUTPUT_HANDLE, 0x4 = ENABLE_VIRTUAL_TERMINAL_PROCESSING
        handle = kernel32.GetStdHandle(-11)
        mode = ctypes.c_uint32()
        if not kernel32.GetConsoleMode(handle, ctypes.byref(mode)):
            return False
        return bool(kernel32.SetConsoleMode(handle, mode.value | 0x4))
    except Exception:
        return False


_LATEX_SYMBOLS = {
    r"\times": "×", r"\cdot": "·", r"\div": "÷", r"\pm": "±", r"\mp": "∓",
    r"\leq": "≤", r"\le": "≤", r"\geq": "≥", r"\ge": "≥", r"\neq": "≠",
    r"\ne": "≠", r"\approx": "≈", r"\equiv": "≡", r"\sim": "∼",
    r"\infty": "∞", r"\sum": "∑", r"\prod": "∏", r"\int": "∫",
    r"\partial": "∂", r"\nabla": "∇", r"\sqrt": "√", r"\angle": "∠",
    r"\in": "∈", r"\notin": "∉", r"\subset": "⊂", r"\subseteq": "⊆",
    r"\supset": "⊃", r"\cup": "∪", r"\cap": "∩", r"\emptyset": "∅",
    r"\varnothing": "∅", r"\forall": "∀", r"\exists": "∃", r"\nexists": "∄",
    r"\therefore": "∴", r"\because": "∵", r"\land": "∧", r"\lor": "∨",
    r"\lnot": "¬", r"\neg": "¬", r"\mid": "|", r"\nmid": "∤",
    r"\Rightarrow": "⇒", r"\Leftrightarrow": "⇔", r"\Leftarrow": "⇐",
    r"\rightarrow": "→", r"\to": "→", r"\leftarrow": "←", r"\mapsto": "↦",
    r"\ldots": "…", r"\dots": "…", r"\cdots": "⋯", r"\vdots": "⋮",
    r"\prime": "′", r"\circ": "∘", r"\bullet": "•", r"\star": "⋆",
    r"\alpha": "α", r"\beta": "β", r"\gamma": "γ", r"\delta": "δ",
    r"\epsilon": "ε", r"\varepsilon": "ε", r"\zeta": "ζ", r"\eta": "η",
    r"\theta": "θ", r"\iota": "ι", r"\kappa": "κ", r"\lambda": "λ",
    r"\mu": "μ", r"\nu": "ν", r"\xi": "ξ", r"\pi": "π", r"\rho": "ρ",
    r"\sigma": "σ", r"\tau": "τ", r"\upsilon": "υ", r"\phi": "φ",
    r"\varphi": "φ", r"\chi": "χ", r"\psi": "ψ", r"\omega": "ω",
    r"\Gamma": "Γ", r"\Delta": "Δ", r"\Theta": "Θ", r"\Lambda": "Λ",
    r"\Xi": "Ξ", r"\Pi": "Π", r"\Sigma": "Σ", r"\Phi": "Φ",
    r"\Psi": "Ψ", r"\Omega": "Ω",
    r"\quad": "  ", r"\qquad": "    ", r"\,": " ", r"\;": " ", r"\!": "",
    r"\{": "{", r"\}": "}", r"\%": "%", r"\$": "$", r"\&": "&", r"\#": "#",
}
# Longest first so \varepsilon wins over \var..., \le doesn't eat \leq.
_SYMBOL_RE = re.compile(
    "|".join(re.escape(k) for k in sorted(_LATEX_SYMBOLS, key=len, reverse=True))
)
_SUB_MAP = str.maketrans("0123456789+-=()aehijklmnoprstuvx",
                         "₀₁₂₃₄₅₆₇₈₉₊₋₌₍₎ₐₑₕᵢⱼₖₗₘₙₒₚᵣₛₜᵤᵥₓ")
_SUP_MAP = str.maketrans("0123456789+-=()abcdefghijklmnoprstuvwxyz",
                         "⁰¹²³⁴⁵⁶⁷⁸⁹⁺⁻⁼⁽⁾ᵃᵇᶜᵈᵉᶠᵍʰᶦʲᵏˡᵐⁿᵒᵖʳˢᵗᵘᵛʷˣʸᶻ")


def _script(body: str, sup: bool) -> str:
    """Unicode sub/superscript when every char maps, else a plain fallback."""
    table = _SUP_MAP if sup else _SUB_MAP
    if body and all(ord(ch) in table for ch in body):
        return body.translate(table)
    # No unicode form (e.g. uppercase has no subscripts) — parenthesise instead
    # of leaving raw LaTeX braces on screen.
    return ("^" if sup else "_") + (f"({body})" if len(body) > 1 else body)


def _brace_arg(text: str, i: int):
    """Read a balanced {...} starting at text[i]; returns (content, next_i)."""
    if i >= len(text) or text[i] != "{":
        return None, i
    depth, j = 0, i
    while j < len(text):
        if text[j] == "{":
            depth += 1
        elif text[j] == "}":
            depth -= 1
            if depth == 0:
                return text[i + 1:j], j + 1
        j += 1
    return text[i + 1:], len(text)  # unclosed (still streaming) — take the rest


def render_latex(src: str) -> str:
    """LaTeX fragment -> readable unicode. Best-effort, never raises."""
    # Wrappers whose braces are just grouping: keep the content, drop the macro.
    for macro in (r"\text", r"\mathrm", r"\mathbf", r"\mathit", r"\mathcal",
                  r"\mathbb", r"\operatorname", r"\textbf", r"\textit", r"\bm"):
        out, i = [], 0
        while i < len(src):
            if src.startswith(macro, i) and not src[i + len(macro):i + len(macro) + 1].isalpha():
                body, i = _brace_arg(src, i + len(macro))
                out.append(body if body is not None else macro)
            else:
                out.append(src[i])
                i += 1
        src = "".join(out)

    # \frac{a}{b} -> a/b, parenthesised when either side is compound.
    for _ in range(6):  # bounded nesting passes
        idx = src.find(r"\frac")
        if idx < 0:
            break
        num, j = _brace_arg(src, idx + 5)
        den, k = _brace_arg(src, j)
        if num is None or den is None:
            break
        wrap = lambda s: s if (len(s) <= 1 or s.isalnum()) else f"({s})"
        src = src[:idx] + f"{wrap(num)}/{wrap(den)}" + src[k:]

    src = re.sub(r"\\sqrt\{([^{}]*)\}", r"√(\1)", src)
    src = re.sub(r"\\(left|right|big|Big|bigg|Bigg)\b", "", src)
    # Bare function names read fine without the backslash.
    src = re.sub(r"\\(max|min|log|ln|exp|sin|cos|tan|lim|gcd|lcm|det|deg|mod)\b",
                 r"\1", src)
    src = _SYMBOL_RE.sub(lambda m: _LATEX_SYMBOLS[m.group(0)], src)

    # Sub/superscripts: braced form first, then the single-character form.
    out, i = [], 0
    while i < len(src):
        ch = src[i]
        if ch in "_^" and i + 1 < len(src):
            sup = ch == "^"
            if src[i + 1] == "{":
                body, i = _brace_arg(src, i + 1)
                out.append(_script(body or "", sup))
                continue
            out.append(_script(src[i + 1], sup))
            i += 2
            continue
        out.append(ch)
        i += 1
    src = "".join(out)

    src = src.replace(r"\\", " ").replace("~", " ")
    return re.sub(r"[ \t]{2,}", " ", src).strip()


_MATH_RE = re.compile(r"\$\$(.+?)\$\$|\$(.+?)\$|\\\[(.+?)\\\]|\\\((.+?)\\\)",
                      re.DOTALL)


def render_line(line: str, color: bool = True) -> str:
    """One line of model output -> terminal-ready text."""
    b, d, i_, c, y, r = (BOLD, DIM, ITAL, CYAN, YELL, RESET) if color else ("",) * 6

    def math_sub(m):
        body = next(g for g in m.groups() if g is not None)
        return f"{c}{render_latex(body)}{r}"

    line = _MATH_RE.sub(math_sub, line)
    # Loose LaTeX outside any $ delimiters — the model emits plenty of it.
    if "\\" in line or re.search(r"[_^]\{", line):
        line = render_latex(line)

    line = re.sub(r"`([^`]+)`", lambda m: f"{y}{m.group(1)}{r}", line)
    line = re.sub(r"\*\*(.+?)\*\*", lambda m: f"{b}{m.group(1)}{r}", line)
    line = re.sub(r"(?<![\w*])\*([^*\n]+?)\*(?![\w*])",
                  lambda m: f"{i_}{m.group(1)}{r}", line)
    heading = re.match(r"^(#{1,6})\s+(.*)$", line)
    if heading:
        line = f"{b}{heading.group(2)}{r}"
    line = re.sub(r"^(\s*)[-*+]\s+", r"\1• ", line)
    return line


def _is_looping(ids: List[int], sizes=(6, 10, 16, 24)) -> bool:
    """True when the tail is an exact back-to-back repeat of a block.

    Catches the greedy-decoding death spiral (the same clause emitted forever)
    without touching text that merely reuses words.
    """
    for n in sizes:
        if len(ids) >= 2 * n and ids[-n:] == ids[-2 * n:-n]:
            return True
    return False


class StreamRenderer:
    """Buffers streamed tokens and prints each finished line rendered.

    Rendering needs whole lines (`**bold**` cannot be styled until the closing
    `**` arrives), so output appears a line at a time rather than a token at a
    time.  Inside ``` fences nothing is rewritten — code stays literal.
    """

    # Lines that are nothing but a display-math delimiter. The model writes
    # \[ on its own line, the equation on the next, \] on a third — so a
    # single-line regex can never see the pair. These switch a mode instead.
    _MATH_OPEN = ("\\[", "$$", "\\begin{equation}", "\\begin{align}",
                  "\\begin{align*}", "\\begin{aligned}")
    _MATH_CLOSE = ("\\]", "$$", "\\end{equation}", "\\end{align}",
                   "\\end{align*}", "\\end{aligned}")

    def __init__(self, enabled: bool = True, color: bool = True):
        self.enabled = enabled
        self.color = color
        self.buf = ""
        self.in_fence = False
        self.in_math = False
        self._live = 0  # visible chars of the current line already echoed

    def feed(self, text: str) -> None:
        if not self.enabled:
            print(text, end="", flush=True)
            return
        self.buf += text
        while "\n" in self.buf:
            line, self.buf = self.buf.split("\n", 1)
            self._rewind()
            self._emit(line)
        # Echo the in-progress tail live, so generation still looks alive
        # instead of stalling until the line ends. _rewind() erases it before
        # the finished line is reprinted with formatting applied. Needs ANSI:
        # without it the erase codes would print as literal garbage, so fall
        # back to plain line-at-a-time output.
        if not self.color:
            return
        tail = self.buf[self._live:]
        if tail:
            try:
                print(tail, end="", flush=True)
                self._live = len(self.buf)
            except UnicodeEncodeError:
                print(tail.encode("ascii", "backslashreplace").decode("ascii"),
                      end="", flush=True)
                self._live = len(self.buf)

    def _rewind(self) -> None:
        """Erase the live-echoed partial line, including any wrapped rows."""
        if not self._live:
            return
        try:
            width = max(20, shutil.get_terminal_size((80, 24)).columns)
        except Exception:
            width = 80
        rows = max(1, (self._live + width - 1) // width)
        # \r to column 0, up (rows-1), then clear everything below.
        print("\r" + (f"\033[{rows - 1}A" if rows > 1 else "") + "\033[J",
              end="", flush=True)
        self._live = 0

    def _emit(self, line: str) -> None:
        stripped = line.strip()
        if stripped.startswith("```"):
            self.in_fence = not self.in_fence
            print(f"{DIM if self.color else ''}{line}{RESET if self.color else ''}",
                  flush=True)
            return
        if self.in_fence:
            print(line, flush=True)
            return

        # Display-math block: swallow the delimiter lines, render what's between
        # them as LaTeX (indented and coloured so it reads as an equation).
        if not self.in_math and stripped in self._MATH_OPEN:
            self.in_math = True
            return
        if self.in_math:
            if stripped in self._MATH_CLOSE:
                self.in_math = False
                return
            try:
                body = render_latex(line)
            except Exception:
                body = line
            c, r = (CYAN, RESET) if self.color else ("", "")
            self._say(f"    {c}{body}{r}" if body.strip() else "")
            return

        try:
            out = render_line(line, self.color)
        except Exception:
            out = line  # never let formatting break a reply
        self._say(out)

    def _say(self, text: str) -> None:
        try:
            print(text, flush=True)
        except UnicodeEncodeError:  # legacy console codepage
            print(text.encode("ascii", "backslashreplace").decode("ascii"), flush=True)

    def flush(self) -> None:
        if self.enabled:
            if self.buf:
                self._rewind()
                self._emit(self.buf)
            self.in_math = False
        else:
            print("", flush=True)
        self.buf = ""
        self._live = 0

@dataclass
class LatentForwardOutput:
    logits: torch.Tensor
    hidden_states: torch.Tensor
    attention_mask: torch.Tensor
    past_key_values: Optional[Any] = None
    latent_count: int = 0

class LatentUpdateGate(nn.Module):
    """Task-Anchored GRU-style residual gate controlling h_{t-1} -> h_t flow."""

    def __init__(self, hidden_size: int, bias: bool = True):
        super().__init__()
        self.gate_proj = nn.Linear(3 * hidden_size, hidden_size, bias=bias)
        self.norm = nn.LayerNorm(hidden_size)

    def forward(
        self,
        h_prev: torch.Tensor,
        h_new: torch.Tensor,
        h_prompt: torch.Tensor,
        return_gate: bool = False,
    ) -> torch.Tensor:
        combined = torch.cat([h_prev, h_new, h_prompt], dim=-1)
        gate = torch.sigmoid(self.gate_proj(combined))
        h_updated = gate * h_new + (1.0 - gate) * h_prev
        h_out = self.norm(h_updated)
        if return_gate:
            return h_out, gate
        return h_out


class HaltHead(nn.Module):
    """PonderNet-style halting head for adaptive computation."""

    def __init__(self, hidden_size: int, intermediate_size: Optional[int] = None, bias: bool = True):
        super().__init__()
        mid = intermediate_size if intermediate_size is not None else max(hidden_size // 4, 64)
        self.fc1 = nn.Linear(hidden_size, mid, bias=bias)
        self.fc2 = nn.Linear(mid, 1, bias=bias)

    def forward(self, h: torch.Tensor) -> torch.Tensor:
        x = F.gelu(self.fc1(h))
        logits = self.fc2(x).squeeze(-1)
        return torch.sigmoid(logits.float()).to(dtype=h.dtype)


class LatentCausalLM(nn.Module):
    """
    Wrapper that bypasses lm_head during internal thought steps.
    """
    def __init__(
        self,
        causal_lm: nn.Module,
        tokenizer: Optional[Any] = None,
        eot_token_id: Optional[int] = None,
    ) -> None:
        super().__init__()
        self.causal_lm = causal_lm
        self.tokenizer = tokenizer
        self.__dict__['backbone'] = self._resolve_backbone(causal_lm)
        self.__dict__['input_embeddings'] = causal_lm.get_input_embeddings()
        out_emb = causal_lm.get_output_embeddings()
        if out_emb is None and hasattr(causal_lm, "lm_head"):
            out_emb = causal_lm.lm_head
        if out_emb is None:
            raise ValueError("Could not resolve model output embeddings / lm_head.")
        self.__dict__['output_embeddings'] = out_emb

        self.eot_token_id = eot_token_id
        if tokenizer is not None:
            self.eot_token_id = tokenizer.convert_tokens_to_ids(EOT_TOKEN)

        emb_dim = int(self.input_embeddings.embedding_dim)
        hidden_size = self._config_hidden_size(causal_lm.config)
        if hidden_size is not None and emb_dim != int(hidden_size):
            raise ValueError(
                f"Input embedding dim ({emb_dim}) must match hidden size ({hidden_size}) "
                "for raw hidden-state reinjection."
            )
        model_type = str(getattr(causal_lm.config, "model_type", "")).lower()
        self.position_id_mode = "absolute" if "qwen" in model_type else "mask"

        # --- LayerNorm Gate ---
        self.__dict__['final_norm'] = self._resolve_final_norm(causal_lm)
        if self.final_norm is not None:
            print(f"LayerNorm gate enabled: using {type(self.final_norm).__name__}", flush=True)

        try:
            self._backbone_accepts_cache_position = (
                "cache_position" in inspect.signature(self.backbone.forward).parameters
            )
        except (TypeError, ValueError):
            self._backbone_accepts_cache_position = False

        # --- Gated Latent Updates & Norms ---
        self.latent_norm = nn.LayerNorm(hidden_size)
        self.latent_update_gate = LatentUpdateGate(hidden_size)
        self.halt_head = HaltHead(hidden_size)
        self.use_gated_latent = False  # default to False, enabled dynamically if weights found in checkpoint

        # Match the backbone's dtype. These modules are constructed in fp32
        # by default while the backbone loads in bf16/fp16; load_state_dict
        # preserves the destination dtype, so cast them explicitly to avoid a
        # dtype mismatch on the first latent step
        # (use_gated_latent=False routes through _latent_norm_gate, which
        # casts explicitly). Without this a gated checkpoint raises:
        #   RuntimeError: mat1 and mat2 must have the same dtype,
        #                 but got BFloat16 and Float
        _dtype = self.input_embeddings.weight.dtype
        self.latent_norm.to(dtype=_dtype)
        self.latent_update_gate.to(dtype=_dtype)
        self.halt_head.to(dtype=_dtype)

        # Optional hook: fn(h, step_idx) -> h on the injection-scaled latent
        # vector right before it enters the backbone. None = no-op (normal
        # chat); left in place for anyone experimenting with the latent path.
        self.latent_intervention = None

    def load_state_dict(self, state_dict, strict=True):
        # --- Fix key prefix mismatch ---
        # Non-FSDP checkpoints save inner.causal_lm.state_dict() which produces
        # bare keys like 'model.layers.0...' and 'lm_head.weight', but the
        # wrapper expects 'causal_lm.model.layers.0...' and 'causal_lm.lm_head.weight'.
        # Detect this and re-prefix automatically.
        has_causal_prefix = any(k.startswith("causal_lm.") for k in state_dict.keys())
        has_bare_model = any(k.startswith("model.") or k.startswith("lm_head.") for k in state_dict.keys())
        if not has_causal_prefix and has_bare_model:
            print("[Wrapper: Detected bare causal_lm keys — adding 'causal_lm.' prefix]", flush=True)
            fixed = {}
            for k, v in state_dict.items():
                # Wrapper-level keys (latent_norm, latent_update_gate, halt_head)
                # should NOT get the prefix — but these won't be present in bare
                # causal_lm checkpoints anyway. Guard just in case.
                if k.startswith(("latent_norm", "latent_update_gate", "halt_head")):
                    fixed[k] = v
                else:
                    fixed[f"causal_lm.{k}"] = v
            state_dict = fixed

        # Migrate old gate_project (2×H) → gate_proj (3×H) with zero-padding
        renames = {k: k.replace("gate_project", "gate_proj") for k in list(state_dict.keys()) if "gate_project" in k}
        for old_k, new_k in renames.items():
            tensor = state_dict.pop(old_k)
            if tensor.dim() == 2:
                h_out, two_h = tensor.shape
                pad = torch.zeros(h_out, two_h // 2, dtype=tensor.dtype, device=tensor.device)
                tensor = torch.cat([tensor, pad], dim=1)
            state_dict[new_k] = tensor
            print(f"[Wrapper: Migrated {old_k} → {new_k}]", flush=True)

        has_gate = any("latent_update_gate" in k for k in state_dict.keys())
        if has_gate:
            self.use_gated_latent = True
            print("[Wrapper: Gated latent update enabled (weights found in checkpoint)]", flush=True)
        else:
            self.use_gated_latent = False
            print("[Wrapper: Gated latent update disabled (running in backward-compatible mode)]", flush=True)
        return super().load_state_dict(state_dict, strict=strict)

    @staticmethod
    def _resolve_backbone(causal_lm: nn.Module) -> nn.Module:
        for attr in ("model", "transformer", "gpt_neox", "backbone"):
            if hasattr(causal_lm, attr):
                return getattr(causal_lm, attr)
        raise ValueError(
            "Could not resolve the transformer backbone. For Llama/Mistral this is model.model."
        )

    @staticmethod
    def _resolve_final_norm(causal_lm: nn.Module) -> Optional[nn.Module]:
        """Find the model's final normalization layer (RMSNorm/LayerNorm)."""
        backbone = None
        for attr in ("model", "transformer", "gpt_neox", "backbone"):
            if hasattr(causal_lm, attr):
                backbone = getattr(causal_lm, attr)
                break
        if backbone is None:
            return None
        for attr in ("norm", "final_layernorm", "ln_f", "final_layer_norm"):
            if hasattr(backbone, attr):
                norm = getattr(backbone, attr)
                if isinstance(norm, nn.Module):
                    return norm
        return None

    def _latent_norm_gate(self, h: torch.Tensor) -> torch.Tensor:
        """
        Approximate the missing trained latent_norm.
        Default LayerNorm forces std=1.0, which is 40x too large for Layer 0
        and causes NaN explosion. We manually scale to match embedding std.
        """
        # Ensure latent_norm is on the correct dtype/device
        self.latent_norm.to(dtype=h.dtype, device=h.device)
        h_norm = self.latent_norm(h)
        
        if hasattr(self, "causal_lm") and hasattr(self.causal_lm, "model"):
            if not hasattr(self, "_cached_emb_std"):
                # Compute on CPU once to avoid DirectML fallback warning and performance hit
                self._cached_emb_std = self.causal_lm.model.embed_tokens.weight.detach().cpu().float().std().item()
            return h_norm * self._cached_emb_std
            
        return h_norm

    @staticmethod
    def _config_hidden_size(config: Any) -> Optional[int]:
        for name in ("hidden_size", "n_embd", "d_model"):
            value = getattr(config, name, None)
            if value is not None:
                return int(value)
        return None

    @property
    def device(self) -> torch.device:
        return self.input_embeddings.weight.device

    def _position_ids_from_mask(
        self,
        attention_mask: torch.Tensor,
        *,
        force_absolute: bool = False,
    ) -> torch.Tensor:
        if force_absolute or self.position_id_mode == "absolute":
            seq_len = attention_mask.shape[-1]
            return torch.arange(seq_len, device=attention_mask.device, dtype=torch.long).unsqueeze(0).expand(
                attention_mask.shape[0],
                seq_len,
            )
        pos = attention_mask.long().cumsum(dim=-1) - 1
        return pos.clamp_min_(0)

    def _backbone_forward(
        self,
        *,
        inputs_embeds: torch.Tensor,
        attention_mask: torch.Tensor,
        position_ids: torch.Tensor,
        past_key_values: Optional[Any] = None,
        use_cache: bool = False,
    ) -> Any:
        kwargs = {
            "inputs_embeds": inputs_embeds,
            "attention_mask": attention_mask,
            "position_ids": position_ids,
            "past_key_values": past_key_values,
            "use_cache": use_cache,
            "return_dict": True,
        }
        if self._backbone_accepts_cache_position:
            kwargs["cache_position"] = position_ids[0]
        return self.backbone(**kwargs)

    def _lm_head(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.output_embeddings(hidden_states)

    def forward_prefix_latents_suffix(
        self,
        *,
        prefix_input_ids: torch.Tensor,
        latent_steps: int | str,
        suffix_input_ids: Optional[torch.Tensor] = None,
        append_eot: bool = False,
        prefix_attention_mask: Optional[torch.Tensor] = None,
        use_cache_for_latents: bool = True,
    ) -> LatentForwardOutput:
        return self._forward_cached(
            prefix_input_ids=prefix_input_ids,
            latent_steps=latent_steps,
            suffix_input_ids=suffix_input_ids,
            append_eot=append_eot,
            prefix_attention_mask=prefix_attention_mask,
        )

    def _forward_cached(
        self,
        *,
        prefix_input_ids: torch.Tensor,
        latent_steps: int | str,
        suffix_input_ids: Optional[torch.Tensor],
        append_eot: bool,
        prefix_attention_mask: Optional[torch.Tensor],
    ) -> LatentForwardOutput:
        prefix_input_ids = prefix_input_ids.to(self.device)
        batch_size, prefix_len = prefix_input_ids.shape
        if prefix_len == 0:
            raise ValueError("prefix_input_ids cannot be empty.")

        if prefix_attention_mask is None:
            attention_mask = torch.ones(
                (batch_size, prefix_len), dtype=torch.long, device=self.device
            )
        else:
            attention_mask = prefix_attention_mask.to(self.device)

        prefix_embeds = self.input_embeddings(prefix_input_ids)
        pos = self._position_ids_from_mask(attention_mask)
        out = self._backbone_forward(
            inputs_embeds=prefix_embeds,
            attention_mask=attention_mask,
            position_ids=pos,
            use_cache=True,
        )

        hidden_pieces = [out.last_hidden_state]
        past = out.past_key_values
        h = out.last_hidden_state[:, -1:, :]
        h_prompt = h.detach()

        is_dynamic = (latent_steps == "dynamic")
        max_steps = 32 if is_dynamic else int(latent_steps)
        actual_steps = 0

        # --- PonderNet-style adaptive halting (trained halt head) -----------
        # p_halt(t) = lam_t * prod_{i<t}(1 - lam_i). lam_t is read from the
        # injection-scaled state at the TOP of step t, before the backbone runs
        # the step.  In dynamic mode we stop once the
        # cumulative halt mass crosses HALT_THRESHOLD (the distribution's
        # median at 0.5); in fixed-step mode the numbers are telemetry only.
        halt_survival = 1.0
        halt_mass = 0.0
        halt_threshold = float(
            getattr(self, "halt_threshold", None)
            or os.environ.get("HALT_THRESHOLD", "0.7")
        )

        for step_idx in range(max_steps):
            # --- LayerNorm Gate: normalize before reinjection ---
            h = self.latent_norm(h) if self.use_gated_latent else self._latent_norm_gate(h)

            # Optional hook: replaces the injected vector (and what the halt
            # head sees) for anyone experimenting with the latent path.
            if self.latent_intervention is not None:
                h = self.latent_intervention(h, step_idx)

            lambda_t = None
            if self.use_gated_latent:
                with torch.no_grad():
                    lambda_t = float(self.halt_head(h).reshape(-1)[0].item())

            one_mask = torch.ones((batch_size, 1), dtype=torch.long, device=self.device)
            attention_mask = torch.cat([attention_mask, one_mask], dim=-1)
            pos = self._position_ids_from_mask(attention_mask)[:, -1:]

            out = self._backbone_forward(
                inputs_embeds=h,
                attention_mask=attention_mask,
                position_ids=pos,
                past_key_values=past,
                use_cache=True,
            )
            h_new = out.last_hidden_state[:, -1:, :]
            h = self.latent_update_gate(h, h_new, h_prompt) if self.use_gated_latent else h_new
            past = out.past_key_values
            hidden_pieces.append(h)
            actual_steps += 1

            # Update halting bookkeeping AFTER the backbone step so the slot
            # is materialized in the KV cache before the halt decision.
            halted = False
            if lambda_t is not None:
                halt_mass += halt_survival * lambda_t
                halt_survival *= (1.0 - lambda_t)
                halted = is_dynamic and halt_mass >= halt_threshold

            # Logit Lens: project h to vocabulary to see "silent" thoughts
            if self.tokenizer is not None:
                with torch.no_grad():
                    step_logits = self._lm_head(h)
                    probs = torch.softmax(step_logits[:, -1, :], dim=-1)
                    top_probs, top_indices = torch.topk(probs, k=5, dim=-1)
                    top_tokens = []
                    for idx, prob in zip(top_indices[0], top_probs[0]):
                        token_text = self.tokenizer.decode([idx.item()])
                        token_repr = repr(token_text).strip("'")
                        top_tokens.append(f"'{token_repr}' ({prob.item()*100:.1f}%)")
                    halt_str = (
                        f" lam={lambda_t:.2f} halt={halt_mass:.2f}"
                        if lambda_t is not None else ""
                    )
                    thought_str = f"  [Thought Step {step_idx+1}]{halt_str}: {', '.join(top_tokens)}"
                    try:
                        print(thought_str, flush=True)
                    except UnicodeEncodeError:
                        print(thought_str.encode('ascii', errors='backslashreplace').decode('ascii'), flush=True)

                    if is_dynamic and not halted:
                        top_token_id = top_indices[0, 0].item()
                        top_token_text = self.tokenizer.decode([top_token_id])
                        # Stop if model predicts <eot> (trained stop token) OR '\n'
                        # '\n' is a secondary "done thinking" signal, used
                        # alongside the <eot> stop token.
                        is_eot = (self.eot_token_id is not None and top_token_id == self.eot_token_id)
                        is_newline = (top_token_text == "\n")
                        if is_eot or is_newline:
                            reason = "<eot>" if is_eot else "'\\n' (done-thinking signal)"
                            print(f"  [Dynamic Stop]: Model predicted {reason} at step {step_idx+1}.", flush=True)
                            break

            # Trained halting signal: the PonderNet halt head says the thought
            # is complete.  Checked outside the logit-lens block so it works
            # even when no tokenizer is attached.
            if halted:
                print(
                    f"  [Halt Head Stop]: cumulative halt mass {halt_mass:.2f} >= "
                    f"{halt_threshold:.2f} at step {step_idx+1}.",
                    flush=True,
                )
                break

        if append_eot:
            eot_ids = torch.full(
                (batch_size, 1),
                int(self.eot_token_id),
                dtype=torch.long,
                device=self.device,
            )
            h = self.input_embeddings(eot_ids)
            one_mask = torch.ones((batch_size, 1), dtype=torch.long, device=self.device)
            attention_mask = torch.cat([attention_mask, one_mask], dim=-1)
            pos = self._position_ids_from_mask(attention_mask)[:, -1:]
            
            out = self._backbone_forward(
                inputs_embeds=h,
                attention_mask=attention_mask,
                position_ids=pos,
                past_key_values=past,
                use_cache=True,
            )
            h = out.last_hidden_state[:, -1:, :]
            past = out.past_key_values
            hidden_pieces.append(h)

        if suffix_input_ids is not None and suffix_input_ids.numel() > 0:
            suffix_input_ids = suffix_input_ids.to(self.device)
            for i in range(suffix_input_ids.shape[1]):
                next_id = suffix_input_ids[:, i:i+1]
                h = self.input_embeddings(next_id)
                one_mask = torch.ones((batch_size, 1), dtype=torch.long, device=self.device)
                attention_mask = torch.cat([attention_mask, one_mask], dim=-1)
                pos = self._position_ids_from_mask(attention_mask)[:, -1:]
                
                out = self._backbone_forward(
                    inputs_embeds=h,
                    attention_mask=attention_mask,
                    position_ids=pos,
                    past_key_values=past,
                    use_cache=True,
                )
                h = out.last_hidden_state[:, -1:, :]
                past = out.past_key_values
                hidden_pieces.append(h)

        full_hidden = torch.cat(hidden_pieces, dim=1)
        return LatentForwardOutput(
            logits=self._lm_head(full_hidden),
            hidden_states=full_hidden,
            attention_mask=attention_mask,
            past_key_values=past,
            latent_count=actual_steps,
        )

    @torch.no_grad()
    def generate_after_latent_thought(
        self,
        *,
        prefix_input_ids: torch.Tensor,
        latent_steps: int,
        max_new_tokens: int,
        eos_token_id: Optional[int] = None,
        temperature: float = 0.0,
        stop_event: Optional[threading.Event] = None,
        tokenizer: Optional[Any] = None,
        repetition_penalty: float = 1.0,
        renderer: Optional["StreamRenderer"] = None,
    ) -> torch.Tensor:
        """Greedy/sample generation after N latent thought steps and an injected <eot>.
        
        Streams tokens to stdout as they are generated.
        If stop_event is set (e.g. by Ctrl+C), generation stops cleanly after the
        current token and control returns to the caller.
        """
        self.eval()
        out = self.forward_prefix_latents_suffix(
            prefix_input_ids=prefix_input_ids,
            latent_steps=latent_steps,
            suffix_input_ids=None,
            append_eot=True,
            use_cache_for_latents=True,
        )
        attention_mask = out.attention_mask
        past = out.past_key_values
        logits = out.logits[:, -1, :]
        generated: List[torch.Tensor] = []
        tok = tokenizer or self.tokenizer
        sink = renderer or StreamRenderer(enabled=False)
        ids_so_far: List[int] = []
        text_so_far = ""
        char_ends: List[int] = []  # text length after each token; maps a char cut to a token cut
        role_flip = False

        # Print prefix for streamed output
        print("Hyper: ", end="", flush=True)

        for _ in range(max_new_tokens):
            # --- Check for user interrupt (Ctrl+C) ---
            if stop_event is not None and stop_event.is_set():
                sink.flush()
                print("[Interrupted]", flush=True)
                break

            # Repetition penalty: divide the logit of any token already used.
            # Greedy decoding has no way out of a degenerate loop on its own —
            # this is what stops "Therefore p+2=2. Therefore p+2=2. ..." runs.
            if repetition_penalty and repetition_penalty != 1.0 and ids_so_far:
                seen = torch.tensor(sorted(set(ids_so_far)), device=logits.device)
                vals = logits.index_select(-1, seen)
                logits = logits.index_copy(
                    -1, seen,
                    torch.where(vals > 0, vals / repetition_penalty,
                                vals * repetition_penalty),
                )

            if temperature and temperature > 0:
                # fp32 for the sample: dividing fp16 logits by a small
                # temperature amplifies them ~10x and the fp16 softmax /
                # multinomial path on DirectML is not reliable there.
                probs = torch.softmax(logits.float() / temperature, dim=-1)
                next_id = torch.multinomial(probs, num_samples=1).to(logits.device)
            else:
                next_id = torch.argmax(logits, dim=-1, keepdim=True)
            generated.append(next_id)
            ids_so_far.append(int(next_id[0, 0].item()))

            # Stream token to console immediately
            if tok is not None:
                token_text = tok.decode([next_id[0, 0].item()], skip_special_tokens=True)
                # Role-marker stop. Past its trained depth the model loses track
                # of whose turn it is and starts writing the USER's next message.
                # Without this that text is streamed AND stored as the
                # assistant's own history, so the next turn sees it and the
                # session degenerates. Cut the reply at the marker instead.
                text_so_far += token_text
                char_ends.append(len(text_so_far))
                cut = -1
                for marker in ("\nuser:", "\nUser:", "\nassistant:",
                               "<|im_start|>", "<|im_end|>"):
                    j = text_so_far.find(marker)
                    if j >= 0 and (cut < 0 or j < cut):
                        cut = j
                # Also catch the bare form. Real output showed " user"
                # on its own line with no colon, which the markers above
                # miss. Only checked at the very start of the reply, so
                # the word "user" mid-sentence is unaffected.
                if cut < 0:
                    _head = text_so_far.lstrip().lower()
                    if re.match(r"user\s*[:\n]", _head) or _head == "user":
                        cut = text_so_far.lower().find("user")
                if cut >= 0:
                    keep = len(text_so_far) - len(token_text)
                    if cut > keep:
                        sink.feed(token_text[:cut - keep])
                    sink.flush()
                    print(f"{DIM}[Stopped: the model started writing a user turn "
                          f"- it is past its trained thinking depth]{RESET}",
                          flush=True)
                    role_flip = True
                    # Drop the marker and everything after it from the RETURNED
                    # ids as well. Truncating only the display would leave the
                    # fabricated turn in `generated`, and the caller decodes
                    # that into conversation history — the exact feedback loop
                    # this stop exists to prevent.
                    n_keep = 0
                    for e in char_ends:
                        if e <= cut:
                            n_keep += 1
                        else:
                            break
                    del generated[n_keep:]
                    break
                sink.feed(token_text)

            if eos_token_id is not None and torch.all(next_id.squeeze(-1) == eos_token_id):
                break

            # Loop breaker: an exact block repeat means the model is stuck.
            # Cheaper and more reliable than tuning the penalty high enough to
            # escape, and it does not distort normal text the way that would.
            if _is_looping(ids_so_far):
                sink.flush()
                print(f"{DIM}[Stopped: repetition loop detected]{RESET}", flush=True)
                break

            emb = self.input_embeddings(next_id.to(self.device))
            one_mask = torch.ones((next_id.shape[0], 1), dtype=torch.long, device=self.device)
            attention_mask = torch.cat([attention_mask, one_mask], dim=-1)
            pos = self._position_ids_from_mask(attention_mask)[:, -1:]
            step_out = self._backbone_forward(
                inputs_embeds=emb,
                attention_mask=attention_mask,
                position_ids=pos,
                past_key_values=past,
                use_cache=True,
            )
            past = step_out.past_key_values
            logits = self._lm_head(step_out.last_hidden_state)[:, -1, :]

        sink.flush()
        # An immediate end-of-turn renders identically to a crash: "Hyper: "
        # and nothing else. Say what actually happened instead.
        _interrupted = stop_event is not None and stop_event.is_set()
        if tok is not None and not role_flip and not _interrupted \
                and not text_so_far.strip():
            print(f"{DIM}[No output — the model ended its turn without writing "
                  f"anything. If this keeps happening, /clear the session.]{RESET}",
                  flush=True)
        print("", flush=True)  # blank line after the streamed response

        if not generated:
            return torch.empty((prefix_input_ids.shape[0], 0), dtype=torch.long, device=self.device)
        return torch.cat(generated, dim=1)

def ensure_hyper_v4_tokens(tokenizer: Any, model: Optional[nn.Module] = None) -> int:
    special_tokens = [BOT_TOKEN, EOT_TOKEN, STEP_TOKEN]
    added = tokenizer.add_special_tokens(
        {"additional_special_tokens": special_tokens}
    )
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token or tokenizer.unk_token
    if model is not None and model.get_input_embeddings().num_embeddings != len(tokenizer):
        model.resize_token_embeddings(len(tokenizer))
    return int(added)

def main():
    # This fires once per generation on DirectML and lands in the middle of
    # the streamed answer, corrupting the visible line.
    warnings.filterwarnings("ignore", message=".*not currently supported on the DML.*")
    if hasattr(sys.stdout, "reconfigure"):
        try:
            sys.stdout.reconfigure(encoding="utf-8", errors="replace")
        except Exception:
            pass
    import argparse
    parser = argparse.ArgumentParser(description="Hyper v1 - Neuralese-Thinking Model")
    parser.add_argument("--device", default=None, choices=["cpu", "cuda", "directml"], help="Force specific device.")
    args = parser.parse_args()

    # 1. Hardware setup
    device_override = args.device
    if device_override == "cpu":
        device = torch.device("cpu")
        print("Forced execution on CPU.", flush=True)
    elif device_override == "cuda":
        device = torch.device("cuda")
        print("Forced execution on CUDA.", flush=True)
    elif device_override == "directml":
        try:
            import torch_directml
            device = torch_directml.device()
            print("Forced execution on DirectML.", flush=True)
        except ImportError:
            print("Error: torch_directml not installed.", file=sys.stderr)
            sys.exit(1)
    else:
        # Default behavior: try directml first, then cuda, then cpu
        try:
            import torch_directml
            device = torch_directml.device()
            print("Using DirectML accelerator.", flush=True)
        except ImportError:
            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
            print(f"Using device: {device}", flush=True)

    # 2. Model files live next to this script
    here = Path(__file__).resolve().parent
    for required in ("model.safetensors", "latent_modules.safetensors",
                     "config.json", "tokenizer.json"):
        if not (here / required).exists():
            print(f"Error: {required} is missing from {here}.", file=sys.stderr)
            sys.exit(1)

    print("Loading tokenizer & config from the model folder...", flush=True)
    tokenizer = AutoTokenizer.from_pretrained(str(here))
    config = AutoConfig.from_pretrained(str(here))

    # Temporarily mock weights initialization to skip CPU random-fill overhead (saves 2 mins)
    import torch.nn.init as init
    old_kaiming = init.kaiming_uniform_
    old_uniform = init.uniform_
    old_normal = init.normal_
    old_constant = init.constant_
    init.kaiming_uniform_ = lambda t, *a, **k: t
    init.uniform_ = lambda t, *a, **k: t
    init.normal_ = lambda t, *a, **k: t
    init.constant_ = lambda t, *a, **k: t

    print("Initializing model architecture from config (instant load)...", flush=True)
    model = AutoModelForCausalLM.from_config(
        config,
        dtype=torch.bfloat16 if device.type == "cpu" else torch.float16
    )
    # Resize embeddings FIRST to match tokenizer additions
    ensure_hyper_v4_tokens(tokenizer, model)

    # Restore original init functions
    init.kaiming_uniform_ = old_kaiming
    init.uniform_ = old_uniform
    init.normal_ = old_normal
    init.constant_ = old_constant

    # Move the empty model to the GPU first to allocate the memory structure
    print(f"Moving model architecture to {device}...", flush=True)
    model = model.to(device)

    print("Wrapping model in LatentCausalLM...", flush=True)
    wrapper = LatentCausalLM(model, tokenizer=tokenizer).to(device)

    # Weights ship as two safetensors files next to this script:
    #   model.safetensors            the fine-tuned Qwen backbone
    #   latent_modules.safetensors   latent_norm / update gate / halt head
    # The wrapper expects backbone keys under a "causal_lm." prefix.
    from safetensors.torch import load_file
    print("Loading weights...", flush=True)
    state = {f"causal_lm.{k}": v for k, v in
             load_file(str(here / "model.safetensors")).items()}
    state.update(load_file(str(here / "latent_modules.safetensors")))
    _missing, _unexpected = wrapper.load_state_dict(state, strict=False)
    if _unexpected:
        print(f"  note: {len(_unexpected)} unexpected key(s): {_unexpected[:3]}",
              flush=True)
    del state

    # Chat loop
    color_ok = enable_ansi()
    renderer = StreamRenderer(enabled=True, color=color_ok)
    print("\n" + "="*50, flush=True)
    print("Hyper v1 — Neuralese-Thinking Model Online", flush=True)
    print("Type 'exit' to quit, /help for commands.", flush=True)
    print("Press Ctrl+C during generation to interrupt and return to prompt.", flush=True)
    print("="*50 + "\n", flush=True)

    # 4 is the measured optimum for this checkpoint (gold-answer CE 8.555 at
    # k=4 vs 8.967 at k=2, and 12.112 with no thinking at all). At k=2 it
    # gets simple arithmetic wrong that it gets right at k=4.
    latent_steps = 4
    max_new_tokens = 1024
    # Greedy by default. Sampling at temperature 0.1 in fp16 on DirectML
    # produced incoherent answers on questions the same model answers
    # correctly at temperature 0. Raise it with /temp if you want variety.
    temperature = 0.0
    repetition_penalty = 1.08
    messages = []
    _generating = False  # True while model is generating a response

    # --- Ctrl+C interrupt handler ---
    # When NOT generating: exits the program (normal behaviour).
    # When generating: sets STOP_GENERATION so the token loop exits cleanly.
    import signal as _signal
    _original_sigint = _signal.getsignal(_signal.SIGINT)

    def _sigint_handler(signum, frame):
        if _generating:
            STOP_GENERATION.set()
        else:
            # Not generating — restore default handler and re-raise to exit
            _signal.signal(_signal.SIGINT, _original_sigint)
            raise KeyboardInterrupt

    _signal.signal(_signal.SIGINT, _sigint_handler)

    while True:
        try:
            user_input = input("You: ").strip()
        except (KeyboardInterrupt, EOFError):
            print("\nExiting.", flush=True)
            break

        if not user_input:
            continue
        if user_input.lower() in ["exit", "quit"]:
            break

        if user_input.lower() in ["/wipe", "/clear"]:
            messages = []
            os.system('cls' if os.name == 'nt' else 'clear')
            print("\n" + "="*50, flush=True)
            print("Hyper v1 — Neuralese-Thinking Model Online", flush=True)
            print("Session cleared and reset. Local memory is fully wiped.", flush=True)
            print("="*50 + "\n", flush=True)
            continue

        if user_input.startswith("/steps "):
            parts = user_input.split()
            if len(parts) > 1:
                val = parts[1].lower()
                if val == "dynamic":
                    latent_steps = "dynamic"
                    print("[System: Latent thinking steps set to dynamic — stops when the trained "
                          "halt head's cumulative halt mass crosses 0.7 (HALT_THRESHOLD env to tune), "
                          "or on a predicted <eot>/'\\n']", flush=True)
                else:
                    try:
                        new_steps = int(val)
                        latent_steps = max(0, new_steps)
                        print(f"[System: Latent thinking steps set to {latent_steps}]", flush=True)
                    except ValueError:
                        print("[System: Invalid format. Use /steps <number> or /steps dynamic]", flush=True)
            continue

        if user_input.startswith("/halt"):
            parts = user_input.split()
            if len(parts) > 1:
                try:
                    wrapper.halt_threshold = min(0.999, max(0.01, float(parts[1])))
                    print(f"[System: Halt threshold set to {wrapper.halt_threshold:.2f}. "
                          f"Higher = more latent steps before the halt head stops it. "
                          f"Only affects /steps dynamic.]", flush=True)
                except ValueError:
                    print("[System: Invalid format. Use /halt <0.01-0.99>]", flush=True)
            else:
                print(f"[System: Halt threshold is {getattr(wrapper, 'halt_threshold', 0.7):.2f}]",
                      flush=True)
            continue

        if user_input.startswith("/render"):
            parts = user_input.split()
            val = parts[1].lower() if len(parts) > 1 else ("off" if renderer.enabled else "on")
            renderer.enabled = (val not in ("off", "0", "false"))
            print(f"[System: Markdown/LaTeX rendering {'ON' if renderer.enabled else 'OFF'}]",
                  flush=True)
            continue

        if user_input.startswith("/penalty"):
            parts = user_input.split()
            if len(parts) > 1:
                try:
                    repetition_penalty = max(1.0, float(parts[1]))
                    print(f"[System: Repetition penalty set to {repetition_penalty:.2f} "
                          f"(1.0 = off)]", flush=True)
                except ValueError:
                    print("[System: Invalid format. Use /penalty <1.0-1.5>]", flush=True)
            else:
                print(f"[System: Repetition penalty is {repetition_penalty:.2f}]", flush=True)
            continue

        if user_input.startswith("/help"):
            print(
                "[System: /steps <n>|dynamic   latent thinking depth (k=2 measured best)\n"
                "         /halt <0-1>          dynamic-mode halt threshold (higher = deeper)\n"
                "         /penalty <1.0-1.5>   repetition penalty\n"
                "         /temp <x>            sampling temperature (0 = greedy)\n"
                "         /max_tokens <n>      generation cap\n"
                "         /render on|off       markdown + LaTeX prettifying\n"
                "         /clear               wipe conversation memory\n"
                "         exit                 quit]", flush=True)
            continue

        if user_input.startswith("/max_tokens "):
            parts = user_input.split()
            if len(parts) > 1:
                try:
                    new_tokens = int(parts[1])
                    max_new_tokens = max(1, new_tokens)
                    print(f"[System: Max generation tokens set to {max_new_tokens}]", flush=True)
                except ValueError:
                    print("[System: Invalid format. Use /max_tokens <number>]", flush=True)
            continue

        if user_input.startswith("/temperature ") or user_input.startswith("/temp "):
            parts = user_input.split()
            if len(parts) > 1:
                try:
                    new_temp = float(parts[1])
                    temperature = max(0.0, new_temp)
                    print(f"[System: Temperature set to {temperature:.2f}]", flush=True)
                except ValueError:
                    print("[System: Invalid format. Use /temperature <value>]", flush=True)
            continue

        # Append user message
        messages.append({"role": "user", "content": user_input})

        # Format input prompt using Qwen's ChatML template (must end with <bot>)
        system_prompt = SYSTEM_PROMPT
        prompt = f"<|im_start|>system\n{system_prompt}\n<|im_end|>\n"
        for msg in messages:
            if msg["role"] == "user":
                prompt += f"<|im_start|>user\n{msg['content']}\n<|im_end|>\n"
            elif msg["role"] == "assistant":
                # History carries the ANSWER text only. Emitting {BOT}{EOT}
                # here puts those two tokens adjacent, which never occurs in
                # training - every trained example has at least one latent slot
                # between them - so it is an out-of-distribution sequence for
                # every past turn. A previous turn's latent thinking has no
                # text form, so omitting it is the honest rendering.
                prompt += f"<|im_start|>assistant\n{msg['content']}\n<|im_end|>\n"
        prompt += f"<|im_start|>assistant\n{BOT_TOKEN}"

        inputs = tokenizer(prompt, return_tensors="pt")
        prefix_ids = inputs.input_ids.to(device)

        steps_label = latent_steps if latent_steps != "dynamic" else "dynamic"
        print(f"Hyper (thinking with {steps_label} steps...):", flush=True)

        # --- Reset interrupt flag and mark as generating ---
        STOP_GENERATION.clear()
        _generating = True
        try:
            output_ids = wrapper.generate_after_latent_thought(
                prefix_input_ids=prefix_ids,
                latent_steps=latent_steps,
                max_new_tokens=max_new_tokens,
                eos_token_id=tokenizer.eos_token_id,
                temperature=temperature,
                stop_event=STOP_GENERATION,
                tokenizer=tokenizer,
                repetition_penalty=repetition_penalty,
                renderer=renderer,
            )
        finally:
            _generating = False

        # Decode full response for conversation history (from generated token ids)
        response = tokenizer.decode(output_ids[0], skip_special_tokens=True).strip()

        # Only add to history if we got a real response (not just an interruption)
        if response:
            messages.append({"role": "assistant", "content": response})

if __name__ == "__main__":
    main()