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

# Model path (relative to the Space root)
MODEL_PATH = "./emoji_deepseekmath-r1"
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
model = AutoModelForCausalLM.from_pretrained(MODEL_PATH).to("cuda" if torch.cuda.is_available() else "cpu")

def emojiq_brainpower():  
    logic = random.randint(40, 100)  
    emoji_confusion = 100 - logic  
    caffeine = random.randint(10, 50)  
    return (  
        f"🧠 {logic}% logic, {emoji_confusion}% 'Wait... is 🍕 a number?!' 🤔, "  
        f"and {caffeine}% caffeine boost ☕🚀"  
    )  

def solve_emoji_math(problem, temperature, max_length, top_p):
    if not problem.strip():
        return "⚠️ Please enter an emoji math problem to solve.", ""
    
    prompt = f"Solve: {problem}\nAnswer:"
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_length=max_length,
            temperature=temperature,
            top_p=top_p,
            do_sample=True
        )
    
    solution = tokenizer.decode(outputs[0], skip_special_tokens=True)
    solved_text = solution.split("Answer:")[-1].strip()
    brainpower = emojiq_brainpower()
    
    return f"✅ Solution: **{solved_text}**", brainpower

with gr.Blocks(title="EmojIQ - Emoji Math Solver") as interface:
    gr.Markdown("# 🔢 EmojIQ: Emoji Math Solver")
    gr.Markdown("### Cracking Emoji Codes, Solving Math with a Smile! 😃➕🎭")
    gr.Markdown("Enter an emoji math problem, and let EmojIQ solve it! 🤓")

    problem_input = gr.Textbox(
        label="📝 Emoji Math Problem",
        placeholder="e.g., 🍎 + 🍎 + 🍎 = 12",
        lines=3
    )

    with gr.Row():
        with gr.Column():
            temperature = gr.Slider(0.0, 1.5, value=0.7, step=0.1, label="Temperature (0 = logical, 1 = creative)")
            max_length = gr.Slider(50, 200, value=100, step=10, label="Max Output Length")
            top_p = gr.Slider(0.0, 1.0, value=0.9, step=0.05, label="Top-p (sampling diversity)")

    solution_output = gr.Textbox(label="Solution", interactive=False)
    brainpower_output = gr.Textbox(label="EmojIQ Brainpower", interactive=False)

    solve_button = gr.Button("🔍 Solve Emoji Math")
    solve_button.click(
        fn=solve_emoji_math,
        inputs=[problem_input, temperature, max_length, top_p],
        outputs=[solution_output, brainpower_output]
    )

    gr.Markdown("---")
    gr.Markdown("🔢 **EmojIQ** - Cracking Emoji Codes, One Equation at a Time! 🚀")

interface.launch()