"""Interactive demo for 3-digit-basic-calc: watch a 1.6M-parameter transformer do
3-digit arithmetic one scratchpad step at a time."""

import re
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_ID = "vmal/3-digit-basic-calc"  # <-- set to your Hub repo id before deploying

model = AutoModelForCausalLM.from_pretrained(MODEL_ID, trust_remote_code=True).eval()
tok = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)

CONTROL = ["<add>", "<sub>", "<mul>", "<div>", "<pos>", "<neg>",
           "<state>", "<step>", "<ans>", "<nan>", "<col>", "<qmul>", "<rem>"]

EXPR_RE = re.compile(r"^\s*-?\d{1,3}\s*[-+*/]\s*-?\d{1,3}\s*=?\s*$")


def pretty_trace(trace: str) -> str:
    """Insert line breaks so each control token starts a new, readable line."""
    text = trace
    for c in CONTROL:
        text = text.replace(c, "\n" + c + " ")
    return text.strip()


def run(expression: str):
    if not EXPR_RE.fullmatch(expression or ""):
        return "—", "Enter something like  842/37  or  213*145  (operands in −999…999)."
    try:
        answer, trace = model.solve(tok, expression.strip(), return_trace=True)
    except Exception as exc:  # noqa: BLE001
        return "error", f"Could not compute: {exc}"
    return answer, pretty_trace(trace)


with gr.Blocks(title="3-digit-basic-calc") as demo:
    gr.Markdown(
        "# 🧮 3-digit-basic-calc\n"
        "A **1.6M-parameter** transformer trained from scratch that does 3-digit "
        "arithmetic by writing out the algorithm — watch its scratchpad below."
    )
    with gr.Row():
        expr = gr.Textbox(label="Expression", value="842/37",
                          placeholder="e.g. 213*145")
        btn = gr.Button("Compute", variant="primary")
    answer = gr.Textbox(label="Answer", interactive=False)
    trace = gr.Textbox(label="The model's scratchpad (its actual reasoning)",
                       lines=16, interactive=False, show_copy_button=True)
    gr.Examples(
        [["842/37"], ["213*145"], ["3/31"], ["-500+500"], ["999*999"], ["12/0"]],
        inputs=expr,
    )
    btn.click(run, inputs=expr, outputs=[answer, trace])
    expr.submit(run, inputs=expr, outputs=[answer, trace])

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