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

# Lightweight CPU-friendly model
model_name = "microsoft/phi-1_5"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
model.generation_config = GenerationConfig.from_pretrained(model_name)
model.generation_config.pad_token_id = model.generation_config.eos_token_id


def solve_math_problem(question):
    messages = [{"role": "user", "content": question}]
    input_tensor = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt")
    input_tensor = input_tensor.to(model.device)

    with torch.no_grad():
        outputs = model.generate(input_tensor, max_new_tokens=150)

    response = tokenizer.decode(outputs[0][input_tensor.shape[1]:], skip_special_tokens=True)
    return response.strip()


with gr.Blocks(css="footer {visibility: hidden}") as demo:
    gr.Markdown("# 🧙‍♂️ Math Wizard AI")
    gr.Markdown("""
        <div style="font-size: 16px; line-height: 1.5">
        Welcome to the <b>Math Wizard</b> – your intelligent assistant for solving math problems of all kinds! <br>
        Ask anything from algebra, calculus, or even about famous mathematicians.<br><br>
        </div>
    """)

    with gr.Tabs():
        with gr.Tab("🧮 General Math"):
            with gr.Row():
                with gr.Column():
                    question_box = gr.Textbox(
                        label="Ask your question here:",
                        placeholder="E.g. What is the derivative of x^2 + 3x + 2?",
                        lines=3
                    )
                    submit_btn = gr.Button("🔍 Solve Now")
                    clear_btn = gr.Button("❌ Clear")

                with gr.Column():
                    answer_box = gr.Textbox(label="📘 Answer from the Wizard", lines=8, interactive=False)
                    copy_btn = gr.Button("📋 Copy Answer")

            submit_btn.click(fn=solve_math_problem, inputs=question_box, outputs=answer_box)
            clear_btn.click(lambda: ("", ""), outputs=[question_box, answer_box])
            copy_btn.click(lambda x: x, inputs=answer_box, outputs=answer_box, show_progress=False)

        with gr.Tab("🧠 Examples & Inspiration"):
            gr.Markdown("""
            <h4>Try asking things like:</h4>
            <ul>
            <li>🧩 Solve the equation x² + 2x - 3 = 0</li>
            <li>📐 What is the Pythagorean theorem?</li>
            <li>📊 Integrate sin(x) from 0 to π</li>
            <li>🧮 Tell me about Euclid</li>
            </ul>
            """)

        with gr.Tab("📚 About"):
            gr.Markdown("""
            <h4>About Math Wizard</h4>
            <p>This assistant is powered by a lightweight AI model that runs smoothly even on CPUs.</p>
            <p>Built with ❤️ using Gradio + HuggingFace Transformers</p>
            <p>Model: <code>microsoft/phi-1_5</code> optimized for reasoning and small footprint.</p>
            """)

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