# =============================================================================
#  Project   : localvm-research
#  File      : experiments/candidate_01/implementation/runtime.py
#  Purpose   : Margin-gated deferred-refinement generation loop (q4 resident
#              base + windowed q8 verification sweeps with rollback)
#  Author    : Simon-Pierre Boucher
#  Contact   : contact@spboucher.ai
#  Created   : 2026-08-12
#  Modified  : 2026-08-12
#  Platform  : macOS / Apple Silicon (arm64) — MLX / Metal
#  License   : All rights reserved (research code)
# =============================================================================
"""Candidate-01 runtime.

Two modes:
  * ``verify-all``  — every window position is checked against the q8 sweep
    (exact q8-greedy output; precision-tiered speculative decoding).
  * ``margin``      — only positions whose base margin < tau are checked
    (approximate; fewer rollbacks — the novel relaxation under test).
"""

from __future__ import annotations

import time
from dataclasses import dataclass, field

import mlx.core as mx
from mlx_lm.models.cache import make_prompt_cache, trim_prompt_cache


@dataclass
class RunStats:
    tokens_out: int = 0
    deferred: int = 0
    sweeps: int = 0
    rollbacks: int = 0
    sweep_time_s: float = 0.0
    gen_time_s: float = 0.0
    sweep_logical_bytes: int = 0
    events: list = field(default_factory=list)


def _greedy_step(model, inp, cache):
    logits = model(inp, cache=cache)[0, -1].astype(mx.float32)
    top2 = mx.topk(logits, 2)
    nxt = mx.argmax(logits)
    mx.eval(top2, nxt)
    v = top2.tolist()
    return int(nxt.item()), abs(v[1] - v[0])


def generate_deferred(base_model, verify_model, tokenizer, prompt_ids: list[int],
                      max_tokens: int, tau: float, window: int,
                      mode: str, q8_bytes: int) -> tuple[list[int], RunStats]:
    """Generate with the resident base model, verifying windows with the q8
    model. Returns (tokens, stats). `verify_model` weights are assumed
    mmap-backed; each sweep accounts `q8_bytes` logical bytes."""
    assert mode in ("verify-all", "margin")
    stats = RunStats()
    tokens: list[int] = []          # accepted generated tokens
    flags: list[bool] = []          # deferred flag per generated token
    verified_upto = 0               # tokens before this index are settled

    cache = make_prompt_cache(base_model)
    # Incremental verify cache: sweeps feed only the unseen suffix, cutting
    # sweep compute from O(context) to O(window). v_pos = tokens of the
    # current sequence already ingested; invariant at sweep: v_pos <= base-1.
    # verify_model may be a StreamingVerifier (layer-streamed weights for
    # models that do not fit beside the resident base).
    is_streaming = hasattr(verify_model, "forward_chunk")
    vcache = make_prompt_cache(verify_model.model if is_streaming else verify_model)
    v_pos = 0
    t0 = time.perf_counter()
    inp = mx.array(prompt_ids)[None]

    def sweep() -> bool:
        """Verify tokens[verified_upto:] with q8; rollback on first flip.
        Returns True if a rollback happened."""
        nonlocal verified_upto, tokens, flags, inp, cache, v_pos
        pending = len(tokens) - verified_upto
        if pending <= 0:
            return False
        s0 = time.perf_counter()
        stats.sweeps += 1
        stats.sweep_logical_bytes += q8_bytes
        full = prompt_ids + tokens
        base_of_window = len(prompt_ids) + verified_upto
        # feed full[v_pos : end-1]; logits rows are absolute v_pos..end-2,
        # predictions for tokens v_pos+1..end-1 ⊇ the pending window
        chunk = full[v_pos : len(full) - 1]
        if is_streaming:
            logits = verify_model.forward_chunk(chunk, vcache)
        else:
            logits = verify_model(mx.array(chunk)[None], cache=vcache)[0]
        sel = logits[base_of_window - 1 - v_pos :].astype(mx.float32)
        v_pos = len(full) - 1
        q8_argmax = mx.argmax(sel, axis=-1)
        mx.eval(q8_argmax)
        q8_next = q8_argmax.tolist()
        flip_at = None
        for j in range(pending):
            i = verified_upto + j
            check = mode == "verify-all" or flags[i]
            if check and q8_next[j] != tokens[i]:
                flip_at = (i, q8_next[j])
                break
        if flip_at is None:
            verified_upto = len(tokens)
            stats.sweep_time_s += time.perf_counter() - s0
            return False
        i, corrected = flip_at
        stats.rollbacks += 1
        # Cache invariant at sweep time: cache = prompt + tokens[:-1] (the
        # newest token was appended but not yet fed). To leave the cache at
        # prompt + tokens[:i], trim (len(tokens)-1 - i) entries.
        trim_prompt_cache(cache, len(tokens) - 1 - i)
        # verify cache must hold a prefix of the corrected sequence:
        # keep exactly prompt + tokens[:i] (= base_next - 1, invariant holds)
        keep = len(prompt_ids) + i
        if v_pos > keep:
            trim_prompt_cache(vcache, v_pos - keep)
            v_pos = keep
        tokens = tokens[:i] + [corrected]
        flags = flags[:i] + [False]
        verified_upto = len(tokens)
        # do NOT feed `corrected` here — the main loop's next _greedy_step
        # feeds tokens[-1], preserving the invariant
        stats.sweep_time_s += time.perf_counter() - s0
        stats.events.append({"rollback_at": i})
        return True

    # Loop until the full budget is generated AND the tail is verified —
    # a rollback in the final window must resume generation, not truncate.
    max_steps = 6 * max_tokens  # safety bound against rollback ping-pong
    steps = 0
    while steps < max_steps:
        steps += 1
        hit_eos = False
        if len(tokens) < max_tokens:
            nxt, margin = _greedy_step(base_model, inp, cache)
            hit_eos = nxt == tokenizer.eos_token_id
            # EOS is a decision too: append it (always checked) so the sweep
            # can veto a premature stop; stripped before returning.
            tokens.append(nxt)
            flags.append(hit_eos or margin < tau)
            if not hit_eos and margin < tau:
                stats.deferred += 1
            inp = mx.array([[nxt]])
            if not hit_eos and len(tokens) - verified_upto < window and len(tokens) < max_tokens:
                continue
        rolled = sweep()
        if rolled:
            if tokens and tokens[-1] == tokenizer.eos_token_id:
                tokens.pop()  # q8 corrected the decision to "stop here"
                break
            inp = mx.array([[tokens[-1]]])  # corrected token, not yet fed
            continue
        if len(tokens) and tokens[-1] == tokenizer.eos_token_id:
            tokens.pop()
            break
        if len(tokens) >= max_tokens:
            break

    stats.tokens_out = len(tokens)
    stats.gen_time_s = time.perf_counter() - t0 - stats.sweep_time_s
    return tokens, stats