"""Reference diffusion sampler extracted from the evaluated A100 experiment.
Fast prefix reuse remains disabled by the loader because numerical checks failed.
"""
import json
import sys
import time
import torch
import torch.nn.functional as F

DEADLINE = float("inf")

def check_time(reserve=0):
    if time.monotonic() >= DEADLINE - reserve:
        raise TimeoutError("Runtime time budget reached")

def log(event, **values):
    print(json.dumps({"event":event, **values}), file=sys.stderr)

def block_mask(length, prefix, device, dtype):
    q = torch.arange(length, device=device)[:, None]
    k = torch.arange(length, device=device)[None, :]
    allowed = (q >= prefix) | (k <= q)
    return torch.zeros((length, length), device=device, dtype=dtype).masked_fill(~allowed, torch.finfo(dtype).min)[None, None]

def rotate_half(x):
    a, b = x.chunk(2, dim=-1)
    return torch.cat((-b, a), -1)


def rotary(x, rot):
    cos, sin = rot
    return x * cos[:, None] + rotate_half(x) * sin[:, None]


class Sampler:
    def __init__(self, model, mask_embedding):
        self.model, self.mask_embedding = model, mask_embedding
        self.fast = False

    def reference(self, ids, masked, p):
        m = self.model.model
        h = m.embed_tokens(ids)
        h = torch.where(masked[..., None], self.mask_embedding.to(h.dtype), h)
        pos = torch.arange(ids.shape[1], device=ids.device)[None]
        rot = m.pos_emb(h, pos)
        attn = block_mask(ids.shape[1], p, ids.device, h.dtype)
        for layer in m.layers:
            h = layer(h, position_embeddings=rot, attention_mask=attn if layer.is_attention_layer else None,
                      position_ids=pos, past_key_values=None)
        return m.embedding_norm(h)

    def prefix_cache(self, prefix):
        m = self.model.model
        h = m.embed_tokens(prefix)
        p = prefix.shape[1]
        pos = torch.arange(p, device=h.device)[None]
        rot = m.pos_emb(h, pos)
        attn = block_mask(p, p, h.device, h.dtype)
        cache = []
        for layer in m.layers:
            normed = layer.operator_norm(h)
            if layer.is_attention_layer:
                a = layer.self_attn
                k = a.k_layernorm(a.k_proj(normed).view(1, p, -1, a.head_dim)).transpose(1, 2)
                v = a.v_proj(normed).view(1, p, -1, a.head_dim).transpose(1, 2)
                cache.append((rotary(k, rot), v))
            else:
                tail = max(1, layer.conv.conv_kernel_size - 1)
                cache.append(normed[:, -tail:])
            h = layer(h, position_embeddings=rot, attention_mask=attn if layer.is_attention_layer else None,
                      position_ids=pos, past_key_values=None)
        return cache

    def canvas_hidden(self, canvas, masked, cache, rot):
        m = self.model.model
        h = m.embed_tokens(canvas)
        h = torch.where(masked[..., None], self.mask_embedding.to(h.dtype), h)
        n = canvas.shape[1]
        for layer, ctx in zip(m.layers, cache):
            x = layer.operator_norm(h)
            if layer.is_attention_layer:
                a = layer.self_attn
                q = rotary(a.q_layernorm(a.q_proj(x).view(1, n, -1, a.head_dim)).transpose(1, 2), rot)
                k = rotary(a.k_layernorm(a.k_proj(x).view(1, n, -1, a.head_dim)).transpose(1, 2), rot)
                v = a.v_proj(x).view(1, n, -1, a.head_dim).transpose(1, 2)
                k, v = torch.cat((ctx[0], k), 2), torch.cat((ctx[1], v), 2)
                y = F.scaled_dot_product_attention(q, k, v, is_causal=False, scale=a.scaling, enable_gqa=True)
                y = a.out_proj(y.transpose(1, 2).reshape(1, n, -1).contiguous())
            else:
                y = layer.conv(torch.cat((ctx, x), 1))[:, -n:]
            h = h + y
            h = h + layer.feed_forward(layer.ffn_norm(h))
        return m.embedding_norm(h)

    def verify_cache(self, probe):
        p = 128
        prefix, canvas = probe[:, :p], probe[:, p:]
        cache = self.prefix_cache(prefix)
        rot = self.model.model.pos_emb(self.model.model.embed_tokens(canvas), torch.arange(p, probe.shape[1], device="cuda")[None])
        results = []
        for rate in (1.0, .5):
            mask = torch.rand_like(canvas, dtype=torch.float) < rate
            fullmask = torch.cat((torch.zeros_like(prefix, dtype=torch.bool), mask), 1)
            ref = self.reference(probe, fullmask, p)[:, p:]
            got = self.canvas_hidden(canvas, mask, cache, rot)
            a = self.model.lm_head(ref).float()
            b = self.model.lm_head(got).float()
            results.append({"rate": rate, "hidden_max_abs": (ref.float()-got.float()).abs().max().item(),
                            "top1_agreement": a.argmax(-1).eq(b.argmax(-1)).float().mean().item(),
                            "mean_logit_abs": (a-b).abs().mean().item()})
        self.fast = all(x["top1_agreement"] >= .98 and x["mean_logit_abs"] <= .03 for x in results)
        log("cache_validation", accepted=self.fast, probes=results)
        return results

    @torch.inference_mode()
    def generate(self, prompt, cfg, max_new_tokens=1024):
        # CUDA synchronization brackets prefix construction, all diffusion steps and EOS handling.
        torch.cuda.synchronize()
        start = time.perf_counter()
        dev = prompt.device
        eos = self.model.config.eos_token_id
        eos = set(eos if isinstance(eos, list) else [eos])
        eos_t = torch.tensor(sorted(eos), device=dev)
        pad = self.model.config.pad_token_id
        generated, calls, canvases = [], 0, 0
        while len(generated) < max_new_tokens:
            check_time(45)
            n = min(cfg["block"], max_new_tokens - len(generated))
            history = prompt.new_tensor([generated[-256:]])
            prefix = torch.cat((prompt, history), 1)
            p = prefix.shape[1]
            canvas = prompt.new_full((1, n), pad)
            masked = torch.ones_like(canvas, dtype=torch.bool)
            if self.fast:
                cache = self.prefix_cache(prefix)
                rot = self.model.model.pos_emb(self.model.model.embed_tokens(canvas), torch.arange(p, p+n, device=dev)[None])
            steps = min(cfg["steps"], n)
            for s in range(steps):
                if self.fast:
                    h = self.canvas_hidden(canvas, masked, cache, rot)[masked]
                else:
                    ids = torch.cat((prefix, canvas), 1)
                    mask = torch.cat((torch.zeros_like(prefix, dtype=torch.bool), masked), 1)
                    h = self.reference(ids, mask, p)[:, p:][masked]
                calls += 1
                logits = self.model.lm_head(h).float()
                logits[:, pad] = -torch.inf
                if cfg["repetition_penalty"] != 1:
                    previous = torch.cat((prefix[0, -128:], canvas[0, ~masked[0]]))
                    unique = previous.unique()
                    values = logits[:, unique]
                    penalty = cfg["repetition_penalty"]
                    logits[:, unique] = torch.where(values >= 0, values / penalty, values * penalty)
                logits /= .9
                if cfg["top_k"]:
                    values, indices = logits.topk(cfg["top_k"], dim=-1)
                    probs = values.softmax(-1)
                    sampled = torch.multinomial(probs, 1)
                    tokens = indices.gather(-1, sampled).squeeze(-1)
                    confidence = probs.gather(-1, sampled).squeeze(-1)
                else:
                    probs = logits.softmax(-1)
                    tokens = torch.multinomial(probs, 1).squeeze(-1)
                    confidence = probs.gather(-1, tokens[:, None]).squeeze(-1)
                remaining_after = n * (steps - s - 1) // steps
                positions = masked[0].nonzero().flatten()
                commit = len(positions) - remaining_after
                chosen = confidence.topk(commit).indices
                dest = positions[chosen]
                canvas[0, dest] = tokens[chosen]
                masked[0, dest] = False
                # Only an EOS with all earlier positions resolved may end a partial canvas.
                hits = torch.isin(canvas[0], eos_t) & ~masked[0]
                if hits.any():
                    first = int(hits.nonzero()[0])
                    if not masked[0, :first].any():
                        break
            block = canvas[0].tolist()
            stop = next((i for i, token in enumerate(block) if token in eos), None)
            canvases += n
            generated.extend(block if stop is None else block[:stop])
            if stop is not None:
                break
        torch.cuda.synchronize()
        seconds = time.perf_counter() - start
        return generated, {"tokens": len(generated), "seconds": seconds, "tokens_per_second": len(generated)/seconds,
                           "forward_calls": calls, "canvas_slots_allocated": canvases}


def health(ids):
    if not ids:
        return {"healthy": False, "distinct4": 0., "max_run": 0}
    run, longest = 1, 1
    for a, b in zip(ids, ids[1:]):
        run = run + 1 if a == b else 1
        longest = max(longest, run)
    grams = [tuple(ids[i:i+4]) for i in range(max(0, len(ids)-3))]
    distinct = len(set(grams)) / len(grams) if grams else 1.
    return {"healthy": longest <= 8 and distinct >= .8, "distinct4": distinct, "max_run": longest}


def encode(tokenizer, text):
    # Use explicit completed-empty thinking section rather than an unfinished <think> prefix.
    rendered = tokenizer.apply_chat_template([{"role": "user", "content": text}], tokenize=False, add_generation_prompt=True)
    assert isinstance(rendered, str)
    if rendered.endswith("<think>"):
        rendered += "</think>"
    return tokenizer(rendered, add_special_tokens=False, return_tensors="pt")["input_ids"].to("cuda")