"""
Nawah-Math-Reasoning — Gradio demo.

The prompt rendering (ChatML + BOS prepend) is IDENTICAL to train_reasoning.py. A 51M model
is very sensitive to format drift, so do not change render_prompt() without changing training.

Runs on ZeroGPU. The model is only ~52M parameters and works on CPU too, but ZeroGPU keeps
responses snappy. `import spaces` must come BEFORE torch so it can patch the CUDA calls.

The model emits <think>…</think> before its answer, so the stream is split live into two
panels: the reasoning trace and the final answer.

Deploy: push this + requirements.txt + README.md to a Gradio Space. The released model is
public, so no token is needed; MODEL_HF_TOKEN is still read for pointing MODEL_ID at a private
checkpoint.
"""

import os
import re
import threading

import spaces  # import BEFORE torch so it can patch CUDA calls
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer

# ── Config ────────────────────────────────────────────────────────────────────

MODEL_ID = os.environ.get("MODEL_ID", "oddadmix/Nawah-Math-Reasoning")
# Unused for the public release; needed only if MODEL_ID is repointed at a private repo.
# HF_TOKEN is reserved by Spaces (a secret set under that name does not reach the container),
# so MODEL_HF_TOKEN is the one to set.
HF_TOKEN = os.environ.get("MODEL_HF_TOKEN") or os.environ.get("HF_TOKEN")

IM_START, IM_END = "<|im_start|>", "<|im_end|>"
THINK_OPEN, THINK_CLOSE = "<think>", "</think>"
MAX_NEW_TOKENS_CAP = 1500

# ── Load (once, at startup) ───────────────────────────────────────────────────

print("[*] token env vars present:",
      [k for k in ("MODEL_HF_TOKEN", "HF_TOKEN") if os.environ.get(k)] or "NONE")
print(f"[*] Loading {MODEL_ID} ...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.bfloat16, token=HF_TOKEN)
model.to("cuda").eval()

CTX = getattr(model.config, "max_position_embeddings", 2048)

_eos = {tokenizer.eos_token_id} if tokenizer.eos_token_id is not None else set()
_im_end = tokenizer.convert_tokens_to_ids(IM_END)
if isinstance(_im_end, int) and _im_end >= 0:
    _eos.add(_im_end)
EOS_IDS = list(_eos) or None
PAD_ID = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id

print(f"[+] {model.num_parameters():,} params | eos_ids={EOS_IDS} | ctx={CTX}")

# ── Prompt rendering — must match train_reasoning.py ──────────────────────────

def render_prompt(question: str) -> str:
    # Single-turn only: the model was trained on one user turn per sample, so a chat
    # history would be out of distribution.
    return f"{IM_START}user\n{question.strip()}{IM_END}\n{IM_START}assistant\n"


STRIP_RE = re.compile(r"<\|im_end\|>|</s>|<pad>|<s>")


def split_stream(text: str):
    """-> (reasoning_so_far, answer_so_far). Handles the partial state mid-stream."""
    text = STRIP_RE.sub("", text)
    if THINK_CLOSE in text:
        reasoning, answer = text.split(THINK_CLOSE, 1)
        return reasoning.replace(THINK_OPEN, "").strip(), answer.strip()
    return text.replace(THINK_OPEN, "").strip(), ""


# ── Generate ──────────────────────────────────────────────────────────────────

@spaces.GPU(duration=60)
def solve(question, max_new_tokens, temperature, repetition_penalty):
    question = (question or "").strip()
    if not question:
        yield "", "", ""
        return

    ids = tokenizer(render_prompt(question), add_special_tokens=False)["input_ids"]
    if tokenizer.bos_token_id is not None:
        ids = [tokenizer.bos_token_id] + ids  # match training's explicit BOS
    input_ids = torch.tensor([ids], device=model.device)

    # skip_special_tokens must stay False — <think>/</think> are real special tokens
    # in this tokenizer, and stripping them would destroy the split.
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=False)

    kwargs = dict(
        input_ids=input_ids,
        attention_mask=torch.ones_like(input_ids),
        max_new_tokens=int(max_new_tokens),
        repetition_penalty=float(repetition_penalty),
        eos_token_id=EOS_IDS,
        pad_token_id=PAD_ID,
        streamer=streamer,
    )
    if temperature and temperature > 0:
        kwargs.update(do_sample=True, temperature=float(temperature), top_p=0.95)
    else:
        kwargs.update(do_sample=False)  # greedy — how the model was evaluated

    threading.Thread(target=model.generate, kwargs=kwargs).start()

    out = ""
    for chunk in streamer:
        out += chunk
        reasoning, answer = split_stream(out)
        yield reasoning, (answer or "…"), out

    reasoning, answer = split_stream(out)
    if not answer:
        answer = "⚠️ لم يُغلق النموذج وسم التفكير — جرّب سؤالًا أقرب لأمثلة التدريب.\n" \
                 "(The model never closed `</think>` — try a question closer to its training distribution.)"
    yield reasoning, answer, out


# ── UI ────────────────────────────────────────────────────────────────────────

DESCRIPTION = """
<div style="text-align:center">
<h1>🧠 Nawah-Math-Reasoning</h1>
<p>نموذج استدلال عربي صغير (~52M بارامتر) يفكّر خطوة بخطوة داخل وسم <code>&lt;think&gt;</code>
ثم يعطي الإجابة النهائية.<br>
A ~52M-parameter Arabic reasoning model that thinks step by step inside
<code>&lt;think&gt;</code> before answering.</p>
<p><i>نموذج صغير بما يكفي ليعمل حتى على المعالج (CPU).<br>
Small enough to run on a CPU — this Space uses ZeroGPU for snappier responses.</i></p>
<p>
<a href="https://huggingface.co/oddadmix/Nawah-Math-Reasoning">Model</a> ·
<a href="https://huggingface.co/oddadmix/Nawah-Math-Reasoning/tree/main/code">Training code</a> ·
<a href="https://huggingface.co/datasets/oddadmix/arabic-math-reasoning-synth">Synthetic dataset</a> ·
<a href="https://huggingface.co/datasets/oddadmix/gsm8k-reasoning-ar">GSM8K-ar dataset</a>
<br><i>Weights, both datasets and the full training code are open — Apache 2.0.</i>
</p>
</div>
"""

NOTE = """
### 📊 النتائج / Results

Number agreement, greedy decoding, on held-out splits — the same rows for every version of the
model, so the numbers are comparable.

| eval set | n | score |
|---|---:|---:|
| GSM8K-ar | 600 | **79.0%** |
| Arabic_Reasoning | 400 | **73.0%** |
| synthetic math | 1000 | **40.4%** |
| synthetic relational (`ضعف`, `نصف`, `أكثر بـ…`) | 400 | **52.2%** |

### ⚠️ حدود النموذج / Limitations
نموذج تجريبي بحجم 52M: يجيد **شكل** الاستدلال العربي ويحلّ مسائل النِّسب والحساب البسيطة،
لكنه **يخطئ في الحساب كثيرًا** — غالبًا خطوات الحل سليمة ثم تقع غلطة في عملية حسابية واحدة
ويكمل النموذج على رقمه الخاطئ. الأسئلة المفتوحة وغير الحسابية خارج نطاقه، والحوار متعدد
الأدوار كذلك.

A 52M proof of concept. It reliably produces the *shape* of Arabic step-by-step reasoning, but
**arithmetic errors are the dominant failure mode**: the derivation is usually structurally
right, one computation is wrong, and the model then stays faithful to its own bad number. The
40.4% and 52.2% above are the honest ceiling on multi-step problems. Single-turn only;
open-ended and non-mathematical questions are out of distribution.
"""

EXAMPLES = [
    "إذا كان لديك 1500 ريال وأنفقت 20% منها على الكتب، فكم تبقى معك؟",
    "في مصنع تم إنتاج 5000 وحدة، وكانت نسبة الوحدات المعيبة 2%، فما عدد الوحدات السليمة؟",
    "في مدرسة بها 500 طالب، إذا كانت نسبة الذكور 55%، فما عدد الطالبات؟",
    "لدى تاجر 240 كيلوغرامًا من الأرز، باع منها 35%، فكم كيلوغرامًا تبقى لديه؟",
    "إذا كان عمر أحمد 12 سنة وعمر أخيه ضعف عمره، فما مجموع عمريهما؟",
    "في حديقة 80 حيوانًا، 25% منها طيور، ونصف الطيور بيضاء. كم عدد الطيور البيضاء؟",
    "جمع سامي 45 صدفة، وجمع أخوه ضعف هذا العدد. كم صدفة جمعا معًا؟",
    "لدى ليلى 60 جنيهًا، ولدى ندى أقل منها بـ 18 جنيهًا. كم معهما معًا؟",
]

with gr.Blocks(title="Nawah-Math-Reasoning") as demo:
    gr.HTML(DESCRIPTION)

    with gr.Row():
        with gr.Column(scale=3):
            question = gr.Textbox(
                label="السؤال / Question", rtl=True, lines=3,
                placeholder="اكتب مسألة حسابية هنا…",
            )
            with gr.Row():
                submit = gr.Button("🧮 حل / Solve", variant="primary")
                clear = gr.Button("مسح / Clear")

            with gr.Accordion("⚙️ إعدادات التوليد / Generation settings", open=False):
                max_new_tokens = gr.Slider(32, MAX_NEW_TOKENS_CAP, value=320, step=8,
                                           label="أقصى عدد توكنز / Max new tokens")
                temperature = gr.Slider(0.0, 1.5, value=0.0, step=0.05,
                                        label="درجة الحرارة / Temperature (0 = greedy, as evaluated)")
                repetition_penalty = gr.Slider(1.0, 1.5, value=1.0, step=0.01,
                                               label="عقوبة التكرار / Repetition penalty")

        with gr.Column(scale=4):
            answer_box = gr.Textbox(label="✅ الإجابة / Answer", rtl=True, lines=3)
            with gr.Accordion("🧠 التفكير / Reasoning trace", open=True):
                reasoning_box = gr.Textbox(label="", rtl=True, lines=12)
            with gr.Accordion("🔍 المخرجات الخام / Raw output", open=False):
                raw_box = gr.Textbox(label="", lines=8)

    gr.Examples(examples=EXAMPLES, inputs=question, label="أمثلة / Examples")
    gr.Markdown(NOTE)

    inputs = [question, max_new_tokens, temperature, repetition_penalty]
    outputs = [reasoning_box, answer_box, raw_box]
    submit.click(solve, inputs=inputs, outputs=outputs)
    question.submit(solve, inputs=inputs, outputs=outputs)
    clear.click(lambda: ("", "", "", ""), outputs=[question] + outputs)

if __name__ == "__main__":
    demo.queue().launch(theme=gr.themes.Soft())