"""
Urdu Education & Reasoning — Gemma-3-4B adapted to Urdu via Adaption AutoScientist.
A compelling results page + live demo, served on HF ZeroGPU.
"""
import os, spaces, gradio as gr, torch
from threading import Thread
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer

MODEL = os.environ.get("MODEL_REPO", "abdullah693/gemma-3-4b-it-urdu-edu-reasoning")

print("loading model on cuda (module level for ZeroGPU)...", flush=True)
tokenizer = AutoTokenizer.from_pretrained(MODEL)
model = AutoModelForCausalLM.from_pretrained(MODEL, device_map="cuda", torch_dtype=torch.bfloat16)
model.eval()
print("ready", flush=True)


@spaces.GPU(duration=120)
def respond(message, history, max_tokens, temperature):
    msgs = []
    for u, a in history:
        msgs.append({"role": "user", "content": u})
        if a: msgs.append({"role": "assistant", "content": a})
    msgs.append({"role": "user", "content": message})
    enc = tokenizer.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt",
                                        return_dict=True).to(model.device)
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    Thread(target=model.generate, kwargs=dict(**enc, streamer=streamer, max_new_tokens=int(max_tokens),
           do_sample=temperature > 0, temperature=temperature if temperature > 0 else None,
           repetition_penalty=1.1, pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id)).start()
    out = ""
    for t in streamer:
        out += t; yield out


CSS = """
@import url('https://fonts.googleapis.com/css2?family=Playfair+Display:wght@700&family=Inter:wght@400;500;600&display=swap');
.gradio-container{max-width:1000px!important;margin:0 auto!important}
#hero{position:relative;border-radius:20px;overflow:hidden;margin-bottom:14px;
  background:linear-gradient(135deg,#01411C 0%,#0a7a43 55%,#01411C 100%);box-shadow:0 14px 40px rgba(1,40,20,.35)}
#hero .in{padding:38px 26px;text-align:center;color:#fff}
#hero h1{font-family:'Playfair Display',serif;font-size:2.4rem;margin:.1em 0;color:#fff;line-height:1.1}
#hero h1 .a{color:#e9c75a}
#hero p{font-family:'Inter',sans-serif;font-size:1.05rem;color:#fff;opacity:1;max-width:680px;margin:10px auto;text-shadow:0 1px 4px rgba(0,0,0,.4)}
#hero .badges{display:flex;gap:8px;justify-content:center;flex-wrap:wrap;margin-top:14px}
#hero .b{font-family:'Inter';font-size:.8rem;color:#06301c;background:#e9c75a;font-weight:600;padding:5px 13px;border-radius:999px}
.sec{font-family:'Inter',sans-serif}
.sec h2{font-family:'Playfair Display',serif;color:#01411C;font-size:1.5rem;margin:.2em 0 .1em}
.kpi{display:flex;gap:10px;flex-wrap:wrap;justify-content:center;margin:6px 0 2px}
.kpi .c{background:#f3f8f4;border:1px solid #d8e6dc;border-radius:14px;padding:12px 18px;text-align:center;min-width:120px}
.kpi .v{font-size:1.7rem;font-weight:700;color:#0a7a43;font-family:'Inter'}
.kpi .l{font-size:.78rem;color:#555}
.note{background:#fbf6e9;border:1px solid #ecdfb8;border-radius:12px;padding:12px 16px;font-family:'Inter';font-size:.92rem;color:#4a4a4a}
#foot{font-family:'Inter';font-size:.8rem;color:#888;text-align:center;margin-top:16px;line-height:1.7}
#foot a{color:#0a7a43;text-decoration:none;font-weight:500}
/* force visible (dark) text in light boxes regardless of light/dark theme */
.note,.note *{color:#3a3a3a!important}
.note b,.note strong{color:#01411C!important;font-weight:700}
.sec h2{color:#01411C!important}
.sec p,.sec p *{color:#26332b!important}
.sec p b,.sec p strong{color:#01411C!important}
.kpi .v{color:#0a7a43!important}.kpi .l{color:#555!important}
#foot,#foot *{color:#7a7a7a!important}#foot a{color:#0a7a43!important}
#hero p,#hero p *,#hero h1{color:#fff!important}#hero h1 .a{color:#e9c75a!important}
"""

HERO = """<div id="hero"><div class="in">
<h1>اردو <span class="a">Education & Reasoning</span></h1>
<p>A Gemma-3-4B model adapted to Urdu by translating English knowledge corpora into Urdu with
<b>Adaption AutoScientist</b>, then benchmarked on UrduMMLU.</p>
<div class="badges"><span class="b">Gemma-3-4B</span><span class="b">Adaption AutoScientist</span>
<span class="b">UrduMMLU 46.2%</span><span class="b">+1.3 vs base</span></div></div></div>"""

KPI = """<div class="kpi">
<div class="c"><div class="v">+5.9</div><div class="l">STEM (pts)</div></div>
<div class="c"><div class="v">+3.6</div><div class="l">Profession</div></div>
<div class="c"><div class="v">+2.7</div><div class="l">Social Sci.</div></div>
<div class="c"><div class="v">46.2%</div><div class="l">UrduMMLU overall</div></div>
<div class="c"><div class="v">+1.3</div><div class="l">vs base overall</div></div></div>"""

VALIDATED = """<div class="sec"><h2>What this validates</h2>
<p>We tested whether adapting English knowledge corpora into Urdu with Adaption AutoScientist improves a 4B
model on a native Urdu benchmark. It does, for knowledge that is language-independent: every such domain
improved and the model exceeded its base overall (44.96% to 46.21%). The effect does not extend to Urdu
literature, which is intrinsic to the language and requires native data rather than translation.</p></div>"""

APPROACH = """<div class="sec"><h2>Method</h2>
<p>Most UrduMMLU subjects test knowledge that is largely language-independent: science, mathematics,
reasoning, and social studies. We assembled about 40,000 examples from open English datasets covering these
subjects, together with native-Urdu instruction and literature data, then used <b>Adaption AutoScientist</b>
and the <b>Adaptive Data</b> pipeline to translate and localise each example into Pakistani Urdu, adding a
reformulated prompt and an English reasoning trace. Gemma-3-4B was supervised-fine-tuned on the result and
evaluated on UrduMMLU, zero-shot.</p></div>"""

FINDING = """<div class="note"><b>Boundary of the method.</b> Cross-lingual adaptation improved the science,
mathematics, reasoning, and social-knowledge domains, but Urdu literature declined by 2.5 points. That
content cannot be produced by translating English sources; improving it requires native Urdu literary data.
This is a limitation of available data, not of the adaptation method.</div>"""

FOOT = """<div id="foot">Model: <a href="https://huggingface.co/abdullah693/gemma-3-4b-it-urdu-edu-reasoning" target="_blank">abdullah693/gemma-3-4b-it-urdu-edu-reasoning</a>
 · Adapted with <b>Adaption AutoScientist</b> · Eval: <a href="https://huggingface.co/datasets/MBZUAI/UrduMMLU" target="_blank">UrduMMLU</a> (Urdu, 0-shot)<br>
<i>Research/educational use — not authoritative for exams or religious rulings.</i></div>"""

EXAMPLES = [
    "اگر ایک ٹرین 60 کلومیٹر فی گھنٹہ کی رفتار سے 2.5 گھنٹے چلے تو کتنا فاصلہ طے کرے گی؟",
    "ایک دکاندار نے 500 روپے کی چیز پر 20 فیصد رعایت دی۔ گاہک کو اب کتنے روپے ادا کرنے ہوں گے؟",
    "مشاعرہ کیا ہوتا ہے اور اردو ثقافت میں اس کی کیا اہمیت ہے؟",
    "تعلیم کسی معاشرے کی ترقی میں کیا کردار ادا کرتی ہے؟ مختصر وضاحت کریں۔",
    "نظامِ شمسی میں کتنے سیارے ہیں اور سب سے بڑا سیارہ کون سا ہے؟",
]

with gr.Blocks(title="Urdu Education & Reasoning", theme=gr.themes.Soft(primary_hue="emerald"), css=CSS) as demo:
    gr.HTML(HERO)
    # ── chat front and centre ──
    gr.HTML('<div class="sec"><h2>Try it — اردو میں سوال پوچھیں</h2></div>')
    with gr.Accordion("⚙️ settings", open=False):
        mt = gr.Slider(64, 512, value=256, step=32, label="Max new tokens")
        tp = gr.Slider(0.0, 1.0, value=0.0, step=0.1, label="Temperature")
    gr.ChatInterface(respond, additional_inputs=[mt, tp], examples=[[e] for e in EXAMPLES], cache_examples=False)
    # ── what we validated ──
    gr.HTML('<div style="margin-top:18px"></div>')
    gr.HTML(VALIDATED)
    gr.HTML(KPI)
    gr.Image("assets/fig_validated.png", show_label=False, container=False)
    # ── results ──
    gr.HTML('<div class="sec" style="margin-top:6px"><h2>Results on UrduMMLU</h2></div>')
    with gr.Row():
        gr.Image("assets/fig_overall.png", show_label=False, container=False)
        gr.Image("assets/fig_domains.png", show_label=False, container=False)
    # ── method ──
    gr.HTML(APPROACH)
    gr.Image("assets/fig_pipeline.png", show_label=False, container=False)
    gr.Image("assets/fig_data.png", show_label=False, container=False)
    gr.HTML(FINDING)
    gr.HTML(FOOT)

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