import os
import spaces
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

BASE_MODEL = "meta-llama/Llama-3.2-3B-Instruct"
ADAPTER = "Minutor/adaption_math_word_problem_sub_2"
TOKEN = os.environ.get("HF_TOKEN")

model = None
tokenizer = None

def load_model():
    global model, tokenizer
    if model is not None:
        return

    print("Loading model...")
    base = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        torch_dtype=torch.bfloat16,
        device_map="cpu",
        token=TOKEN,
        trust_remote_code=True
    )
    model = PeftModel.from_pretrained(base, ADAPTER, token=TOKEN)
    model.eval()

    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, token=TOKEN)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    print("Model loaded!")

SYSTEM_PROMPT = (
    "You are an expert mathematical reasoning assistant. "
    "Always provide a detailed step-by-step solution.\n"
    "Structure your answer like this:\n"
    "1. Understand: What is being asked? What numbers do we have?\n"
    "2. Plan: Should we work forwards or backwards?\n"
    "3. Solution: Show each calculation step with equations and explain why.\n"
    "4. Verify: Plug the answer back in to check.\n"
    "Final Answer: State the answer clearly.\n"
    "Avoid unnecessary variables. Only treat a quantity as zero if the problem explicitly says it is 0, zero, none, or empty - otherwise treat it as an unknown."
)

@spaces.GPU
def solve(question: str) -> str:
    if not question or not question.strip():
        return "Please enter a math word problem."

    load_model()
    model.to("cuda")

    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": question.strip()}
    ]

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

    with torch.inference_mode():
        outputs = model.generate(
            **inputs,
            max_new_tokens=1024,
            do_sample=False,
            pad_token_id=tokenizer.eos_token_id
        )

    answer = tokenizer.decode(
        outputs[0][inputs["input_ids"].shape[1]:],
        skip_special_tokens=True
    )
    return answer.strip()

demo = gr.Interface(
    fn=solve,
    inputs=gr.Textbox(
        lines=4,
        placeholder="Enter a math word problem...",
        label="Math Word Problem"
    ),
    outputs=gr.Textbox(lines=14, label="Step-by-step Solution"),
    title="Math Word Problem Solver (Adaption LoRA)",
    description="Fine-tuned Llama-3.2-3B-Instruct • AutoScientist Challenge (Math & Code)",
    examples=[
        ["A farmer has chickens and cows. Altogether the animals have 50 heads and 140 legs. How many chickens and how many cows does the farmer have?"],
        ["The sum of three consecutive even integers is 150. What is the largest of these three integers?"],
        ["A laptop originally costs $800. It is first discounted by 20%, then an additional 8% tax is applied on the discounted price. What is the final price?"],
        ["I have some $5 notes and $10 notes. Altogether I have 18 notes and their total value is $130. How many $5 notes and how many $10 notes do I have?"]
    ]
)

demo.launch()