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

# ── Load Model ─────────────────────────────────────────────
print("Loading model...")
MODEL_ID = "newtechdevng/qwen-math-tutor"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype = torch.float16,
    device_map  = "cpu",
)
model.eval()
print("✅ Model loaded!")

# ── Inference Function ─────────────────────────────────────
def solve(question, history):
    messages = [
        {"role": "system", "content": "You are a helpful math tutor. Solve problems step by step, showing all working clearly."},
    ]
    for msg in history:
        if msg["role"] in ("user", "assistant"):
            messages.append({"role": msg["role"], "content": msg["content"]})
    messages.append({"role": "user", "content": question})

    text = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = tokenizer(text, return_tensors="pt")

    with torch.no_grad():
        out = model.generate(
            **inputs,
            max_new_tokens = 512,
            temperature    = 0.1,
            do_sample      = True,
            pad_token_id   = tokenizer.eos_token_id,
        )
    return tokenizer.decode(
        out[0][inputs["input_ids"].shape[1]:],
        skip_special_tokens=True
    )

# ── Response Handler ───────────────────────────────────────
def respond(question, history):
    if not question.strip():
        return "", history
    answer = solve(question, history)
    history = history + [
        {"role": "user",      "content": question},
        {"role": "assistant", "content": answer},
    ]
    return "", history

# ── UI ─────────────────────────────────────────────────────
with gr.Blocks(title="Math Tutor AI") as demo:
    gr.Markdown("""
    # 🧮 Math Tutor AI
    ### Powered by Qwen2.5-Math — Fine-tuned for NCERT & competitive math
    ---
    """)

    # Store history as state
    history_state = gr.State([])

    chatbot = gr.Chatbot(height=450)

    with gr.Row():
        question = gr.Textbox(
            placeholder = "Type your math question here...",
            label       = "Your Question",
            scale       = 4,
        )
        submit = gr.Button("Solve →", variant="primary", scale=1)

    clear = gr.Button("🗑️ Clear Chat")

    gr.Examples(
        label    = "Try these examples:",
        examples = [
            "Solve: 2x² - 7x + 3 = 0",
            "Find the area of a triangle with base 12 cm and height 8 cm.",
            "If sin θ = 3/5, find cos θ and tan θ.",
            "A train travels 360 km in 4 hours. Find its speed in m/s.",
            "Find the compound interest on ₹5000 at 10% per annum for 2 years.",
            "Prove that √2 is irrational.",
        ],
        inputs = question,
    )

    gr.Markdown("*Model: Qwen2.5-Math-1.5B fine-tuned on 10K math problems*")

    # ── Event Handlers ─────────────────────────────────────
    def respond_and_display(question, history):
        if not question.strip():
            return "", history, [[m["content"] for m in history[i:i+2]] 
                                  for i in range(0, len(history), 2)]
        answer = solve(question, history)
        history = history + [
            {"role": "user",      "content": question},
            {"role": "assistant", "content": answer},
        ]
        # Convert to tuple pairs for display
        display = []
        for i in range(0, len(history), 2):
            u = history[i]["content"]   if i   < len(history) else ""
            a = history[i+1]["content"] if i+1 < len(history) else ""
            display.append((u, a))
        return "", history, display

    def clear_all():
        return "", [], []

    submit.click(
        respond_and_display,
        [question, history_state],
        [question, history_state, chatbot]
    )
    question.submit(
        respond_and_display,
        [question, history_state],
        [question, history_state, chatbot]
    )
    clear.click(clear_all, outputs=[question, history_state, chatbot])

demo.launch(theme=gr.themes.Soft())