"""
Estella Explainer — HuggingFace Spaces / Lightning AI Demo
A math reading coach for grades 3–8, based on Usable Math (usablemath.org)

Developed by Sai Gattupalli, PhD
College of Education, University of Massachusetts Amherst
Principal Scientist, Society and AI Research Group

Model: sgattup/EstellaExplainerLLM (Gemma 4 E4B fine-tuned)
"""

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

MODEL_ID = "sgattup/EstellaExplainerLLM"

SYSTEM_PROMPT = """You are Estella Explainer, a math reading and language coach for young learners in grades 3 through 8. Your motto is: "My job is to explain math questions clearly so you know what you are supposed to do to solve the problem."

Your rules:
- Give only ONE hint per problem.
- Do NOT solve the problem or show the answer.
- Do NOT show calculation steps or compute anything.
- Help the student understand what the problem is ASKING.
- Use very simple words and short sentences (Flesch Reading Ease 90-100).
- Identify what is known and what is being asked.
- Use phrases like "We know...", "We are looking for...", or "Think about...".
- Be friendly and encouraging."""

EXAMPLES = [
    "A rectangle has a length of 12 cm and a width of 5 cm. What is the area?",
    "Sam had 3/4 of a pizza. He ate 1/4. How much is left?",
    "A store sells apples for $0.75 each. Jake buys 6 apples. How much does he spend?",
    "Round 4,567 to the nearest hundred.",
    "There are 48 students put into equal groups of 6. How many groups are there?",
    "A bag has 5 red marbles and 3 blue marbles. What fraction of the marbles are red?",
    "A train leaves at 10:15 AM and arrives at 1:45 PM. How long is the trip?",
    "The perimeter of a square is 36 inches. What is the length of one side?",
]

print(f"Loading {MODEL_ID}...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float16,
    device_map="auto"
)
model.eval()
print("Estella is ready.")


def ask_estella(message: str, history: list) -> str:
    """Chat function — history is kept for UI display only; Estella gives one hint per problem."""
    if not message.strip():
        return "Please type a math word problem and I'll help explain what it's asking!"

    messages = [
        {"role": "user", "content": SYSTEM_PROMPT + f"\n\n### Math Problem:\n{message.strip()}\n\n### Estella's Hint:"}
    ]
    text = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        tokenize=False
    )
    inputs = tokenizer(text=[text], return_tensors="pt")
    inputs = {k: v.to(model.device) for k, v in inputs.items()}
    input_len = inputs["input_ids"].shape[1]

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=150,
            temperature=0.7,
            do_sample=True,
            repetition_penalty=1.1,
        )

    hint = tokenizer.decode(outputs[0][input_len:], skip_special_tokens=True).strip()
    return hint


# Gradio ChatInterface
demo = gr.ChatInterface(
    fn=ask_estella,
    title="📚 Estella Explainer — Math Reading Coach",
    description="""**"My job is to explain math questions clearly so you know what you are supposed to do to solve the problem."**

Type any math word problem (grades 3–8) and Estella will help you understand what it's asking — without solving it.

*Developed by [Sai Gattupalli, PhD](https://huggingface.co/sgattup) · UMass Amherst College of Education · [Society and AI Research Group](https://societyandai.org) · Based on [Usable Math](https://usablemath.org)*""",
    examples=EXAMPLES,
    cache_examples=False,
    type="messages",
)

demo.launch()