"""DSSD demo: exact self-speculative decoding with pruned exit heads and context copying.

Local UI work without a GPU:  DSSD_DEMO_BACKEND=tiny python app.py
"""

import gradio as gr

from demo.backend import MAX_NEW_TOKENS, MODELS, generate_updates, load_models
from demo.examples import EXAMPLES
from demo.render import legend_html

def gpu_seconds(model, domain_label, prompt, max_new, compare):
    # ZeroGPU refuses a call whose requested duration exceeds the visitor's remaining daily quota
    # (5 min on free accounts), so ask for what the request needs: engine setup + both decodes.
    return int(30 + int(max_new) * (0.25 if compare else 0.12))


try:  # ZeroGPU on Hugging Face Spaces
    import spaces
    GPU = spaces.GPU(duration=gpu_seconds)
except ImportError:
    def GPU(f):
        return f

ENGINES = load_models()
DOMAINS = {"Web text": "web", "Math (few-shot)": "gsm8k", "Python code": "code"}

HEADER = """
<div class="hero">
  <h1>DSSD <span class="hero-sub">exact self-speculative decoding</span></h1>
  <p>A tiny exit head on layer 4 of the frozen model, plus tokens copied from the context, propose a
  small tree of draft tokens. The full model checks the whole tree in one forward pass and keeps the
  longest correct branch. The text is <b>exactly</b> what greedy decoding would write, only faster.</p>
  <div class="badges"><span>✓ identical to greedy decoding</span><span>no draft model</span>
  <span>≈ 280 MB of heads</span></div>
</div>
"""

HOW = """
<div class="how">
  <div class="how-step"><b>1 · Draft</b><br>Copy n-grams found earlier in the context<br>+ chain of the
  layer-4 head<br>+ its top-4 alternatives at the root</div>
  <div class="how-arrow">→</div>
  <div class="how-step"><b>2 · Verify</b><br>One forward pass of the full model over the whole tree
  (16 rows, row-invariant kernels)</div>
  <div class="how-arrow">→</div>
  <div class="how-step"><b>3 · Commit</b><br>Longest branch that matches the model's own greedy choice,
  plus one token from the model</div>
</div>
<p class="how-note">The exit heads predict over a small, layer-specific vocabulary (≈ 3k tokens at layer 4)
plus a <i>reject</i> class, about 50× cheaper than a full projection. Speedups depend on the GPU and on
whether CUDA graphs are used. The paper's measurements (AMD MI210, greedy, batch 1) are 2.50× / 1.72× / 1.46× on
web text, GSM8K and HumanEval for Qwen3-8B-Base, and 2.51× / 1.80× / 1.57× for Llama-3.1-8B.</p>
"""

CSS = """
.hero h1 { margin: 0 0 .3rem 0; font-size: 1.9rem; }
.hero-sub { font-weight: 400; opacity: .65; font-size: 1.1rem; margin-left: .4rem; }
.hero p { margin: .2rem 0 .6rem 0; max-width: 60rem; }
.badges span { display: inline-block; padding: .15rem .6rem; margin: 0 .4rem .3rem 0; border-radius: 999px;
  font-size: .82rem; background: var(--block-background-fill); border: 1px solid var(--border-color-primary); }
:root { --tok-copy: #cfe4fb; --tok-head: #ffe1b0; --tok-star: #f3d1ec; --tok-draft: #d9f2d0; }
.dark { --tok-copy: #1f4b78; --tok-head: #7a5413; --tok-star: #6b2f60; --tok-draft: #2f5a2a; }
.tok { border-radius: 3px; padding: 0 1px; }
.tok-copy { background: var(--tok-copy); }
.tok-head { background: var(--tok-head); }
.tok-star { background: var(--tok-star); box-shadow: inset 0 -2px 0 #b3479b; }
.tok-draft { background: var(--tok-draft); }
.tok-model { background: transparent; border: 1px dashed var(--border-color-primary); }
.legend { display: flex; flex-wrap: wrap; gap: .9rem; font-size: .88rem; margin: .2rem 0; }
.legend-item { display: inline-flex; align-items: center; gap: .35rem; }
.swatch { font-size: .8rem; padding: 0 .3rem; }
.dssd-text { white-space: pre-wrap; overflow-wrap: anywhere; font-family: var(--font-mono); font-size: .9rem; line-height: 1.55;
  min-height: 14rem; max-height: 30rem; overflow-y: auto; padding: .8rem; border-radius: 8px;
  background: var(--block-background-fill); border: 1px solid var(--border-color-primary); }
.cursor { display: inline-block; width: .5em; height: 1em; background: var(--body-text-color);
  vertical-align: text-bottom; animation: blink 1s steps(1) infinite; opacity: .6; }
@keyframes blink { 50% { opacity: 0; } }
.panel-title { font-weight: 600; margin-bottom: .2rem; }
.stats { display: flex; flex-wrap: wrap; gap: .6rem; margin: .4rem 0 .6rem 0; }
.stat { flex: 1 1 8rem; padding: .6rem .8rem; border-radius: 8px; background: var(--block-background-fill);
  border: 1px solid var(--border-color-primary); }
.stat-value { font-size: 1.35rem; font-weight: 700; }
.stat-label { font-size: .78rem; opacity: .7; }
.stat-main .stat-value { font-size: 1.9rem; color: var(--color-accent); }
.stat-ok .stat-value { color: #1a7f37; } .stat-bad .stat-value { color: #cf222e; }
.share-bar { display: flex; height: .7rem; border-radius: 999px; overflow: hidden;
  border: 1px solid var(--border-color-primary); }
.share-bar .tok-model { background: var(--border-color-primary); border: 0; }
.share-labels { font-size: .8rem; opacity: .75; margin-top: .25rem; }
.progress { font-size: .9rem; opacity: .75; padding: .4rem 0; }
.replay-clock { font-size: .78rem; opacity: .7; margin-left: .4rem; }
.how { display: flex; flex-wrap: wrap; align-items: center; gap: .5rem; }
.how-step { flex: 1 1 12rem; padding: .7rem; border-radius: 8px; background: var(--block-background-fill);
  border: 1px solid var(--border-color-primary); font-size: .88rem; }
.how-arrow { font-size: 1.4rem; opacity: .5; }
.how-note { font-size: .85rem; opacity: .8; margin-top: .6rem; }
@media (max-width: 700px) { .how-arrow { display: none; } }
"""

# Client-side replay of both generations side by side, at the speed measured on the GPU.
HEAD = """
<script>
window.dssdReplay = function () {
  const holder = document.querySelector('#replay-holder .replay-data');
  if (!holder) return;
  const data = JSON.parse(holder.dataset.replay);
  const el = {dssd: document.querySelector('#dssd-out .dssd-text'), ar: document.querySelector('#ar-out .dssd-text')};
  const esc = s => s.replace(/&/g, '&amp;').replace(/</g, '&lt;').replace(/>/g, '&gt;');
  for (const k in el) if (el[k]) el[k].innerHTML = '';
  const idx = {dssd: 0, ar: 0}, start = performance.now();
  function frame(now) {
    const t = (now - start) / 1000; let more = false;
    for (const k of ['dssd', 'ar']) {
      const arr = data[k] || [], box = el[k]; if (!box) continue;
      let html = '';
      while (idx[k] < arr.length && arr[idx[k]][0] <= t) {
        const [, txt, src] = arr[idx[k]++];
        html += (k === 'dssd' && src !== 'model') ? `<span class="tok tok-${src}">${esc(txt)}</span>` : esc(txt);
      }
      if (html) box.insertAdjacentHTML('beforeend', html);
      if (idx[k] < arr.length) more = true;
      else if (arr.length && !box.dataset.done) {
        box.dataset.done = 1;
        box.insertAdjacentHTML('beforeend', `<span class="replay-clock">⏱ ${arr[arr.length - 1][0].toFixed(2)} s</span>`);
      }
    }
    if (more) requestAnimationFrame(frame);
    else for (const k in el) if (el[k]) delete el[k].dataset.done;
  }
  requestAnimationFrame(frame);
};
</script>
"""


@GPU
def run(model, domain_label, prompt, max_new, compare):
    try:
        for u in generate_updates(ENGINES[model], prompt, DOMAINS[domain_label], max_new, compare):
            yield u.dssd, u.ar, u.stats, u.replay, gr.update(visible=u.done and compare)
    except ValueError as e:
        raise gr.Error(str(e))


def example(domain_label, i):
    ex = EXAMPLES[DOMAINS[domain_label]]
    i = (i + 1) % len(ex)
    return ex[i], i


with gr.Blocks(title="DSSD: exact self-speculative decoding") as demo:
    gr.HTML(HEADER)
    with gr.Row(equal_height=False):
        with gr.Column(scale=2):
            prompt = gr.Textbox(label="Prompt (plain-text completion: these are base models)", lines=6, max_lines=14,
                                value=EXAMPLES["web"][0])
            with gr.Row():
                ex_btn = gr.Button("Another example", size="sm", variant="secondary")
                ex_idx = gr.State(0)
        with gr.Column(scale=1):
            model = gr.Radio(list(ENGINES), value=next(iter(ENGINES)), label="Model")
            domain = gr.Radio(list(DOMAINS), value="Web text", label="Domain (calibrated settings and stop rule)")
            max_new = gr.Slider(16, MAX_NEW_TOKENS, value=160, step=16, label="Max new tokens")
            compare = gr.Checkbox(value=True, label="Also run plain autoregressive decoding and compare")
            with gr.Row():
                go = gr.Button("Generate", variant="primary")
                stop = gr.Button("Stop", variant="stop")
    gr.HTML(legend_html())
    with gr.Row(equal_height=True):
        with gr.Column():
            gr.HTML('<div class="panel-title">DSSD</div>')
            out_dssd = gr.HTML(elem_id="dssd-out")
        with gr.Column():
            gr.HTML('<div class="panel-title">Autoregressive (greedy)</div>')
            out_ar = gr.HTML(elem_id="ar-out")
    stats = gr.HTML()
    replay_btn = gr.Button("▶ Replay both side by side at the measured speed", visible=False, size="sm")
    replay = gr.HTML(elem_id="replay-holder", visible=True)
    with gr.Accordion("How it works", open=False):
        gr.HTML(HOW)
    gr.Markdown(f"Heads: {', '.join(f'[{r}](https://huggingface.co/{r})' for r in MODELS.values())}. "
                "Llama-3.1-8B heads: built with Llama.", elem_classes="footer")

    gen = go.click(run, [model, domain, prompt, max_new, compare], [out_dssd, out_ar, stats, replay, replay_btn])
    stop.click(None, cancels=[gen])
    replay_btn.click(None, js="() => { window.dssdReplay(); }")
    domain.change(lambda d: (EXAMPLES[DOMAINS[d]][0], 0), domain, [prompt, ex_idx])
    ex_btn.click(example, [domain, ex_idx], [prompt, ex_idx])

demo.queue(default_concurrency_limit=1, max_size=16)

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Soft(primary_hue="indigo", neutral_hue="slate"), css=CSS, head=HEAD)