import streamlit as st
from transformers import pipeline
import torch

# Professional Academic Layout
st.set_page_config(page_title="B.Sc. Math Engine", page_icon="📐")

st.title("📐 B.Sc. Mathematics Engine")
st.write("Advanced Proof & Logic Solver | SUST Dept. of Mathematics")

@st.cache_resource
def load_engine():
    # SmolLM2-1.7B: The best middle-ground for logic vs speed on free CPU
    model_id = "HuggingFaceTB/SmolLM2-1.7B-Instruct"
    return pipeline(
        "text-generation", 
        model=model_id, 
        device=-1, # Force CPU usage
        torch_dtype=torch.float32
    )

pipe = load_engine()

# Strict Academic System Prompt
system_instruction = (
    "You are a Senior Mathematics Professor at a top-tier university. "
    "Follow these strict logical rules for all solutions:\n\n"
    
    "1. STRUCTURE: Always start with 'Definitions & Prerequisites', followed by 'Step-by-Step Proof', and end with 'Conclusion'.\n"
    
    "2. GROUP THEORY: To prove a subgroup H is normal, you must show gH = Hg for all g in G. "
    "If the index [G:H]=2, use the partition G = H U (G-H) to show left and right cosets are identical.\n"
    
    "3. NUMBER THEORY: For divisibility (e.g., n^k - n), use Fermat's Little Theorem or Induction. "
    "For congruences ax ≡ b (mod n), always calculate d = gcd(a, n) first. If d does not divide b, state 'No Solution'.\n"
    
    "4. REAL ANALYSIS: For limits, use the epsilon-delta or epsilon-N definition. "
    "For convergence, state which test you are using (Ratio, Root, or Comparison).\n"
    
    "5. COMPLEX ANALYSIS: To check differentiability, verify Cauchy-Riemann equations (ux = vy and uy = -vx). "
    "For integrals, identify poles and use the Residue Theorem.\n"
    
    "6. FORMATTING: Use LaTeX for EVERY mathematical symbol. Use '$$' for centered equations and '$' for inline math. "
    "Never say 'it is obvious'; explain every logical jump."
)
)

user_query = st.chat_input("Enter your problem...")

if user_query:
    with st.chat_message("user"):
        st.write(user_query)

    with st.chat_message("assistant"):
        with st.spinner("Processing logic... (Usually takes 45-90s on CPU)"):
            # Format specifically for SmolLM2
            prompt = f"<|im_start|>system\n{system_instruction}<|im_end|>\n<|im_start|>user\n{user_query}<|im_end|>\n<|im_start|>assistant\n"
            
            # do_sample=False ensures the math is consistent (Temperature 0)
            result = pipe(prompt, max_new_tokens=600, do_sample=False)
            response = result[0]['generated_text'].split("<|im_start|>assistant\n")[-1]
            
            st.markdown(response)

st.sidebar.info("Model: SmolLM2-1.7B | Mode: Exact Logic")