"""min-spark-preview — a research demo for Meiosis, PICO release 01.

A 5.76M-parameter looped-hybrid language model. The signature mechanic: one
weight-shared body block runs `K` times per token (the "effort"). More passes
sharpen grammar; fewer favor commonsense.

Runs on ZeroGPU (free for the creator; visitors consume their own quota). The
model is tiny so cold-start weight streaming is near-instant.
"""
from __future__ import annotations

import spaces  # MUST precede any torch / CUDA-touching import (ZeroGPU hijack)
import torch
import gradio as gr

from loader import load_model, load_tokenizer, generate

# ── load at module scope, .to("cuda") eagerly so the hijack packs weights ────
# "cuda" as a STRING (never an int device id) — ZeroGPU re-allocs device ids.
print("Loading Meiosis on ZeroGPU ...")
DEVICE = "cuda"
MODEL = load_model(DEVICE)
TOKENIZER = load_tokenizer()
_PARAMS = sum(p.numel() for p in MODEL.parameters())
print(f"  {round(_PARAMS/1e6, 2)}M params ready on {DEVICE}")

# effort label → loop count. The Radio VALUES are what the API/MCP sees.
EFFORT_CHOICES = [("Low (commonsense)", 2),
                  ("Med (grammar)", 3),
                  ("High (deeper grammar)", 4)]


def _estimate_duration(prompt, effort, max_new, temperature, top_k):
    # Tiny model on an RTX PRO 6000: cold-start weight stream (~1-2s) + ~0.5s
    # per 128 tokens. Declare the realistic worst case; cap polite.
    return min(60, 8 + int(max_new) * 0.1)


@spaces.GPU(duration=_estimate_duration)
def run_generate(prompt: str, effort: int, max_new: int,
                 temperature: float, top_k: int):
    """Generate text from a prompt. `effort` is the loop count K (2/3/4): more
    passes sharpen grammar, fewer favor commonsense. Streams token-by-token."""
    k = int(effort) if effort in (2, 3, 4) else 3
    if not str(prompt).strip():
        yield "", "<div class='status'>Type a prompt to start.</div>"
        return
    text = ""
    count = 0
    tps = 0.0
    for chunk, count, tps in generate(
            MODEL, TOKENIZER, prompt, loops=k, max_new=int(max_new),
            temperature=float(temperature), top_k=int(top_k), device=DEVICE):
        text += chunk
        yield text, f"<div class='status'>K={k} · {count} tok · {tps:.1f} tok/s</div>"
    yield text, f"<div class='status'>K={k} · {count} tok · {tps:.1f} tok/s · done</div>"


CSS = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600&family=IBM+Plex+Mono:wght@400;500&display=swap');

:root{
  --bg:#0A1128; --surface:#101A3A; --surface2:#16224A;
  --hair:#22315E; --text:#E8EDF7; --muted:#7D8DB0;
  --blue:#3E7BFA; --blue-bright:#5E96FF; --blue-dim:#2A4A8F;
  --r:10px;
}
*{box-sizing:border-box;}
body, .gradio-container{ background:var(--bg) !important; }
.gradio-container{
  max-width:820px !important; margin:0 auto !important;
  padding:32px 20px 40px !important;
  font-family:'Inter',system-ui,sans-serif !important; color:var(--text) !important;
}

/* hide Gradio chrome: footer, settings, API/MCP links */
footer, .gradio-container footer, #footer{ display:none !important; }

.head{ margin-bottom:20px; }
.title{ font-weight:600; font-size:26px; letter-spacing:-0.01em; color:var(--text); line-height:1.1; }
.title em{ font-style:normal; color:var(--blue-bright); }
.sub{ font-family:'IBM Plex Mono',monospace; font-size:11px; letter-spacing:0.16em;
      text-transform:uppercase; color:var(--muted); margin-top:5px; }

.field-h{ font-family:'IBM Plex Mono',monospace; font-size:10.5px; letter-spacing:0.16em;
          text-transform:uppercase; color:var(--muted); margin:18px 0 8px; }

/* prompt */
.gradio-container textarea{
  background:var(--surface) !important; border:1px solid var(--hair) !important;
  border-radius:var(--r) !important; color:var(--text) !important;
  font-family:'Inter' !important; font-size:15px !important; line-height:1.5 !important;
}
.gradio-container textarea:focus{ border-color:var(--blue) !important; box-shadow:none !important; }

/* effort radio — even pill row */
#effort, #effort .wrap, #effort .wrap > label{ width:100%; }
#effort .wrap{ display:flex !important; gap:8px !important; }
#effort label{
  flex:1 1 0; display:flex !important; align-items:center !important; justify-content:center !important;
  background:var(--surface) !important; border:1px solid var(--hair) !important;
  border-radius:var(--r) !important; padding:11px 8px !important; cursor:pointer;
  font-size:13px !important; color:var(--muted) !important; text-align:center;
  white-space:nowrap; transition:border-color .15s, background .15s, color .15s;
}
#effort label:hover{ border-color:#33508f !important; }
#effort label.selected{
  background:var(--surface2) !important; border-color:var(--blue) !important;
  color:var(--blue-bright) !important; box-shadow:0 0 0 1px var(--blue) inset !important;
}
#effort input{ display:none !important; }

/* sliders */
.gradio-container .form label, .label-wrap label{
  color:var(--muted) !important; font-size:12px !important; letter-spacing:0.02em !important;
}
input[type=range]{ accent-color:var(--blue); }

/* generate button */
button#gen{
  background:var(--blue) !important; color:#fff !important; border:none !important;
  border-radius:var(--r) !important; font-weight:600 !important; font-size:14px !important;
  margin-top:6px;
}
button#gen:hover{ background:var(--blue-bright) !important; }

/* output */
#out textarea, #out{
  background:var(--surface) !important; border:1px solid var(--hair) !important;
  border-radius:var(--r) !important; color:var(--text) !important;
  font-size:15.5px !important; line-height:1.6 !important;
}
.status{ font-family:'IBM Plex Mono',monospace; font-size:11.5px; color:var(--muted);
         margin-top:8px; min-height:16px; }

.foot{ margin-top:26px; border-top:1px solid var(--hair); padding-top:12px;
       font-family:'IBM Plex Mono',monospace; font-size:10.5px; color:#54648c;
       display:flex; justify-content:space-between; flex-wrap:wrap; gap:8px; }
.foot a{ color:var(--muted); text-decoration:none; border-bottom:1px dotted var(--hair); }
.foot a:hover{ color:var(--text); }
"""

HEADER_HTML = """
<div class="head">
  <div class="title">min-spark <em>preview</em></div>
  <div class="sub">Minima Labs</div>
</div>
"""

with gr.Blocks(css=CSS, title="min-spark · Meiosis preview") as demo:
    gr.HTML(HEADER_HTML)

    gr.HTML('<div class="field-h">Effort — loops per token</div>')
    effort = gr.Radio(choices=EFFORT_CHOICES, value=3, elem_id="effort",
                      show_label=False, container=False)

    gr.HTML('<div class="field-h">Prompt</div>')
    prompt = gr.Textbox(value="", placeholder="Once upon a time",
                        lines=2, show_label=False, container=False)

    with gr.Row():
        max_new = gr.Slider(8, 256, value=128, step=8, label="Max tokens")
        temperature = gr.Slider(0.1, 1.6, value=0.8, step=0.05, label="Temperature")
        top_k = gr.Slider(0, 200, value=50, step=5, label="Top-k (0 = off)")  # infer.py default

    gen_btn = gr.Button("Generate", elem_id="gen", variant="primary")

    out = gr.Textbox(value="", elem_id="out", show_label=False, container=False,
                     lines=10, interactive=False, autoscroll=True)
    status = gr.HTML(value='<div class="status"></div>')

    gen_btn.click(
        fn=run_generate,
        inputs=[prompt, effort, max_new, temperature, top_k],
        outputs=[out, status],
        api_name="run_generate",
    )

if __name__ == "__main__":
    demo.queue().launch()