import random
from typing import Dict, List, Set

import spaces
import torch
import torch.nn.functional as F
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_ID = "WhirlwindAI/Arithmetic-SLM"

IM_START = "[IM_START]"
IM_END = "[IM_END]"
NO_THINK = "/no think"
CTX_LEN = 2048

STOP_STRINGS = [IM_END, IM_START]

# ---------------------------------------------------------------------------
# Load model + tokenizer once, at module scope, moved eagerly to CUDA so
# ZeroGPU can pack the weights and stream them into VRAM on the first call.
# The model uses the pure-torch attention backend (config:
# attention_backend="torch", torch_fallback=True) so no flash kernels are
# needed at runtime.
# ---------------------------------------------------------------------------
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    dtype=torch.bfloat16,
    trust_remote_code=True,
).to("cuda")
model.eval()


# ---------------------------------------------------------------------------
# Sampling helpers — ported 1:1 from the model repo's inference.py so output
# matches the authors' reference path exactly.
# ---------------------------------------------------------------------------
def apply_repetition_penalty(logits, generated_ids, penalty):
    if penalty is None or penalty == 1.0:
        return logits
    for tid in set(generated_ids):
        if tid < 0 or tid >= logits.numel():
            continue
        if logits[tid] > 0:
            logits[tid] = logits[tid] / penalty
        else:
            logits[tid] = logits[tid] * penalty
    return logits


def apply_frequency_presence_penalty(logits, generated_ids, frequency_penalty, presence_penalty):
    if not generated_ids:
        return logits
    if frequency_penalty == 0.0 and presence_penalty == 0.0:
        return logits
    counts: Dict[int, int] = {}
    for tid in generated_ids:
        counts[tid] = counts.get(tid, 0) + 1
    for tid, count in counts.items():
        if tid < 0 or tid >= logits.numel():
            continue
        if frequency_penalty:
            logits[tid] -= frequency_penalty * count
        if presence_penalty:
            logits[tid] -= presence_penalty
    return logits


def get_banned_ngram_tokens(generated_ids, no_repeat_ngram_size) -> Set[int]:
    n = no_repeat_ngram_size
    banned: Set[int] = set()
    if n <= 0:
        return banned
    if len(generated_ids) + 1 < n:
        return banned
    prefix_len = n - 1
    current_prefix = tuple(generated_ids[-prefix_len:])
    ngram_map: Dict[tuple, Set[int]] = {}
    for i in range(len(generated_ids) - n + 1):
        prefix = tuple(generated_ids[i:i + prefix_len])
        next_token = generated_ids[i + prefix_len]
        ngram_map.setdefault(prefix, set()).add(next_token)
    banned.update(ngram_map.get(current_prefix, set()))
    return banned


def apply_no_repeat_ngram(logits, generated_ids, no_repeat_ngram_size):
    if no_repeat_ngram_size <= 0:
        return logits
    banned = get_banned_ngram_tokens(generated_ids, no_repeat_ngram_size)
    for tid in banned:
        if 0 <= tid < logits.numel():
            logits[tid] = -float("inf")
    return logits


def apply_top_k(logits, top_k):
    if top_k is None or top_k <= 0:
        return logits
    top_k = min(top_k, logits.size(-1))
    values, _ = torch.topk(logits, top_k)
    cutoff = values[-1]
    logits[logits < cutoff] = -float("inf")
    return logits


def apply_top_p(logits, top_p):
    if top_p is None or top_p >= 1.0:
        return logits
    if top_p <= 0:
        return logits
    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
    sorted_probs = F.softmax(sorted_logits, dim=-1)
    cumulative = torch.cumsum(sorted_probs, dim=-1)
    remove = cumulative > top_p
    remove[1:] = remove[:-1].clone()
    remove[0] = False
    indices_to_remove = sorted_indices[remove]
    logits[indices_to_remove] = -float("inf")
    return logits


def sample_next_token(logits, generated_ids, temperature, top_k, top_p,
                      repetition_penalty, frequency_penalty, no_repeat_ngram_size):
    logits = logits.float().clone()
    logits = apply_repetition_penalty(logits, generated_ids, repetition_penalty)
    logits = apply_frequency_presence_penalty(logits, generated_ids, frequency_penalty, 0.0)
    logits = apply_no_repeat_ngram(logits, generated_ids, no_repeat_ngram_size)

    if temperature <= 0:
        return int(torch.argmax(logits).item())

    logits = logits / temperature
    logits = apply_top_k(logits, top_k)
    logits = apply_top_p(logits, top_p)

    probs = F.softmax(logits, dim=-1)
    if torch.isnan(probs).any() or torch.isinf(probs).any() or probs.sum() <= 0:
        return int(torch.argmax(logits).item())
    return int(torch.multinomial(probs, num_samples=1).item())


def build_stop_sequences(stop_strings) -> List[List[int]]:
    out = []
    for s in stop_strings:
        ids = tokenizer.encode(s, add_special_tokens=False)
        if ids:
            out.append(ids)
    return out


def endswith_sequence(ids, suffix) -> bool:
    if not suffix or len(ids) < len(suffix):
        return False
    return ids[-len(suffix):] == suffix


def strip_after_stop_text(text, stop_strings) -> str:
    best = None
    for s in stop_strings:
        if not s:
            continue
        pos = text.find(s)
        if pos != -1 and (best is None or pos < best):
            best = pos
    return text if best is None else text[:best]


def build_prompt(expression: str, use_think_format: bool) -> str:
    if use_think_format:
        return (
            f"{IM_START}user\n"
            f"{expression} {NO_THINK}"
            f"{IM_END}\n"
            f"{IM_START}assistant\n"
            "<think>\n</think>\n"
        )
    return expression


@spaces.GPU(duration=30)
def solve(
    expression: str,
    use_think_format: bool = True,
    temperature: float = 0.5,
    top_k: int = 40,
    top_p: float = 0.95,
    max_new_tokens: int = 48,
    seed: int = -1,
) -> str:
    """Solve an arithmetic expression with the Arithmetic-SLM model.

    Args:
        expression: An arithmetic expression ending in '=', e.g. '(10 + 28) * 3 ='.
        use_think_format: Use the production [IM_START]/[IM_END] chat template with a '/no think' tag.
        temperature: Sampling temperature (lower = more deterministic).
        top_k: Top-k sampling cutoff.
        top_p: Nucleus (top-p) sampling cutoff.
        max_new_tokens: Maximum number of tokens to generate.
        seed: RNG seed; -1 for random.

    Returns:
        The model's completion of the expression (typically the solved result).
    """
    expression = (expression or "").strip()
    if not expression:
        return "Please enter an arithmetic expression, e.g. '(10 + 28) * 3 ='."

    if seed is not None and int(seed) >= 0:
        random.seed(int(seed))
        torch.manual_seed(int(seed))
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(int(seed))

    repetition_penalty = 1.05
    frequency_penalty = 0.10
    no_repeat_ngram_size = 4
    min_new_tokens = 1

    prompt = build_prompt(expression, use_think_format)

    encoded = tokenizer(prompt, return_tensors="pt", add_special_tokens=False)
    encoded.pop("token_type_ids", None)
    idx = encoded["input_ids"].to("cuda")

    stop_sequences = build_stop_sequences(STOP_STRINGS)
    eos_id = tokenizer.eos_token_id
    generated: List[int] = []

    with torch.no_grad():
        for step in range(int(max_new_tokens)):
            idx_cond = idx[:, -CTX_LEN:]
            out = model(input_ids=idx_cond)
            logits = out.logits[:, -1, :][0]

            if step < min_new_tokens:
                if eos_id is not None and 0 <= eos_id < logits.numel():
                    logits[eos_id] = -float("inf")
                for seq in stop_sequences:
                    if len(seq) == 1 and 0 <= seq[0] < logits.numel():
                        logits[seq[0]] = -float("inf")

            next_id = sample_next_token(
                logits, generated, float(temperature), int(top_k), float(top_p),
                repetition_penalty, frequency_penalty, no_repeat_ngram_size,
            )
            idx = torch.cat(
                [idx, torch.tensor([[next_id]], dtype=torch.long, device=idx.device)], dim=1
            )
            generated.append(next_id)

            if step >= min_new_tokens:
                if eos_id is not None and next_id == eos_id:
                    break
                full_ids = idx[0].tolist()
                if any(endswith_sequence(full_ids, seq) for seq in stop_sequences):
                    break

    full_text = tokenizer.decode(idx[0].tolist(), skip_special_tokens=False)

    if use_think_format:
        # Show the completion after the prompt, cleaned of control markers.
        if full_text.startswith(prompt):
            completion = full_text[len(prompt):]
        else:
            pos = full_text.rfind(prompt)
            completion = full_text[pos + len(prompt):] if pos != -1 else full_text
        completion = strip_after_stop_text(completion, STOP_STRINGS)
        return completion.strip()

    # Raw mode: return the full continued expression, matching the reference
    # inference script's behavior exactly (strip only at [IM_END]/[IM_START]).
    completion = strip_after_stop_text(full_text, STOP_STRINGS)
    return completion.strip()


CSS = """
#col-container { max-width: 820px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # 🧮 Arithmetic-SLM

            A tiny (31.7M parameter) specialized language model that **completes arithmetic
            expressions** — it learned to do math token by token, not with a calculator.
            Handles operator precedence, parentheses, and decimals.

            Enter an expression ending in `=` and let the model finish it.

            [Model card](https://huggingface.co/WhirlwindAI/Arithmetic-SLM)
            """
        )
        with gr.Row():
            expression = gr.Textbox(
                label="Arithmetic expression",
                placeholder="(10 + 28) * 3 =",
                value="(10 + 28) * 3 =",
                scale=4,
            )
            run = gr.Button("Solve", variant="primary", scale=1)

        output = gr.Textbox(label="Model output", lines=3)

        with gr.Accordion("Advanced settings", open=False):
            use_think_format = gr.Checkbox(
                label="Use production /no think chat template (recommended)",
                value=True,
                info=(
                    "Wraps the input in the [IM_START]/[IM_END] template with a "
                    "<think> block — the format the model was trained for, and "
                    "the most reliable. Uncheck for raw free-continuation mode."
                ),
            )
            temperature = gr.Slider(0.0, 1.0, value=0.5, step=0.05, label="Temperature")
            top_k = gr.Slider(0, 100, value=40, step=1, label="Top-k")
            top_p = gr.Slider(0.1, 1.0, value=0.95, step=0.01, label="Top-p")
            max_new_tokens = gr.Slider(8, 128, value=48, step=1, label="Max new tokens")
            seed = gr.Number(value=-1, precision=0, label="Seed (-1 = random)")

        gr.Examples(
            examples=[
                ["59 + 45 ="],
                ["16 + 4 * 3 ="],
                ["(16 / 4) + 44 ="],
                ["3 * 9 + 12 / 1 ="],
                ["(132 / 12) + (46 - 15) ="],
                ["0.5 * 0.5 ="],
                ["8 * 5 + 4 / 4 ="],
                ["(85 - 45) + 56 ="],
            ],
            inputs=[expression],
            outputs=output,
            fn=solve,
            cache_examples=True,
            cache_mode="lazy",
        )

    inputs = [expression, use_think_format, temperature, top_k, top_p, max_new_tokens, seed]
    run.click(solve, inputs=inputs, outputs=output, api_name="solve")
    expression.submit(solve, inputs=inputs, outputs=output, api_name=False)

if __name__ == "__main__":
    demo.launch(mcp_server=True)