from collections import deque
from pathlib import Path

import gradio as gr
import numpy as np
import onnxruntime as ort


ROOT = Path(__file__).parent
MODEL_PATH = ROOT / "model_int8.onnx"
VOCAB_PATH = ROOT / "vocab.json"
TEMPERATURE = 0.9
TOP_K = 32
TOKENS_PER_UPDATE = 24


def load_vocab():
    import json

    vocab = json.loads(VOCAB_PATH.read_text(encoding="utf-8"))
    return vocab["stoi"], {int(key): value for key, value in vocab["itos"].items()}


session = ort.InferenceSession(str(MODEL_PATH), providers=["CPUExecutionProvider"])
stoi, itos = load_vocab()
h_shape = session.get_inputs()[1].shape
num_layers = int(h_shape[0])
hidden_size = int(h_shape[2])


def sample(logits, rng):
    scores = np.asarray(logits, dtype=np.float64).reshape(-1) / TEMPERATURE
    eos_id = stoi.get("EOS")
    if eos_id is not None:
        scores[eos_id] -= 1.0
    k = min(TOP_K, scores.size)
    ids = np.argpartition(scores, -k)[-k:]
    values = scores[ids]
    probabilities = np.exp(values - values.max())
    probabilities /= probabilities.sum()
    return int(rng.choice(ids, p=probabilities))


def step(token_id, h, c, rng):
    outputs = session.run(
        ["logits", "h_out", "c_out"],
        {
            "input": np.array([[token_id]], dtype=np.int64),
            "h": h,
            "c": c,
        },
    )
    return sample(outputs[0], rng), outputs[1], outputs[2]


def generate():
    rng = np.random.default_rng()
    h = np.zeros((num_layers, 1, hidden_size), dtype=np.float32)
    c = np.zeros_like(h)
    prompt = ["BOS", "BPM_120", "GRID_64", "BAR", "POS_0"]
    prompt_ids = [stoi[token] for token in prompt if token in stoi]
    current_id = prompt_ids[0] if prompt_ids else 0

    for token_id in prompt_ids[1:]:
        _, h, c = step(current_id, h, c, rng)
        current_id = token_id

    history = deque((itos[token_id] for token_id in prompt_ids), maxlen=600)
    while True:
        for _ in range(TOKENS_PER_UPDATE):
            current_id, h, c = step(current_id, h, c, rng)
            history.append(itos.get(current_id, "<?>"))
        yield " ".join(history)


with gr.Blocks(title="NanoMaestro CPU Demo") as demo:
    gr.Markdown(
        "**For faster, fully local inference, try the "
        "[Transformers.js browser demo](https://huggingface.co/spaces/utkucoban/NanoMaestro-Realtime).**"
    )
    gr.Markdown("# NanoMaestro CPU Symbolic Generator")
    output = gr.Textbox(label="Continuous symbolic events", lines=18, max_lines=18)
    with gr.Row():
        start = gr.Button("Start", variant="primary")
        stop = gr.Button("Stop")

    generation = start.click(generate, outputs=output)
    stop.click(fn=None, cancels=[generation])


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