import streamlit as st
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

class MathTutor:
    def __init__(self):
        self.model_id = "analist/deepseek-math-tutor-cpu"
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)
        self.model = AutoModelForCausalLM.from_pretrained(
            self.model_id,
            torch_dtype=torch.float32,
            low_cpu_mem_usage=True,
            device_map="cpu"
        )
        
    def get_response(self, question):
        prompt = f"""Below is an instruction that describes a task, paired with an input that provides further context.
Write a response that appropriately completes the request.
Before answering, think carefully about the question and create a step-by-step chain of thoughts to ensure a logical and accurate response.
Your goal is to teach maths a beginner so make it friendly and accessible. Break down your chain of thoughts as for him/her to understand.

### Instruction:
You are a maths expert with advanced knowledge in pedagogy, arithmetics, geometry, analysis, calculus.
Please answer the following questions.

### Question:
{question}

### Response:
<think>"""
        
        inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024)
        outputs = self.model.generate(
            **inputs, 
            max_new_tokens=1200,
            temperature=0.7,
            do_sample=True
        )
        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)

def main():
    st.title("🧮 Friendly Math Tutor")
    st.write("Ask me any math question! I'll help you understand step by step.")
    
    tutor = MathTutor()
    
    question = st.text_area("Your math question:", height=100)
    
    if st.button("Get Help"):
        if question:
            with st.spinner("Thinking..."):
                response = tutor.get_response(question)
                explanation = response.split("### Response:")[1]
                st.markdown(explanation)
        else:
            st.warning("Please enter a question!")
            
    st.divider()
    st.markdown("""
    Example questions:
    - How do I solve quadratic equations?
    - Explain the concept of derivatives
    - Help me understand trigonometry ratios
    """)

if __name__ == "__main__":
    main()