src/streamlit_app.py

import altair as alt
import streamlit as st
import pandas as pd
import os
import zipfile
from PIL import Image
import sympy as sp
import time
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
import tempfile
import pickle
import pickletools

# Page config
st.set_page_config(page_title="🎓 Smart Math Teacher", layout="centered")

# Custom CSS
st.markdown("""
    <style>
    .fun-title { font-size: 40px; color: #ff3399; text-align: center; font-family: 'Comic Sans MS', cursive; }
    .question-box { border: 4px dotted #ffcc00; padding: 20px; border-radius: 20px; background-color: #fff7e6; font-size: 20px; }
    .stButton > button { font-size: 18px; background-color: #00cc99; color: white; border-radius: 10px; padding: 10px; }
    .stTextInput > div > input { font-size: 18px; }
    </style>
""", unsafe_allow_html=True)

# Welcome title
st.markdown("<div class='fun-title'>🧠✨ Welcome to the Smart Math Teacher! ✨🧠</div>", unsafe_allow_html=True)

# Initialize session state for temp directory
if 'temp_dir' not in st.session_state:
    st.session_state.temp_dir = tempfile.mkdtemp()

# Age group setup 
age_groups = {
    "4-6 Age Group": {"dataset": "Datase_of_4-6_Age_Group.xlsx", "zip_file": "Image_for_group_4-6.zip", "image_folder": "Image_for_group_4-6"},
    "7-9 Age Group": {"dataset": "Datase_of_7-9_Age_Group.xlsx", "zip_file": "Image_for_group_7-9.zip", "image_folder": "Image_for_group_7-9"},
    "13-15 Age Group": {"dataset": "Datase_of_13-15_Age_Group.xlsx", "zip_file": "Image_for_group_13-15.zip", "image_folder": "Image_for_group_13-15"},
}

selected_age_group = st.selectbox("🧒 Select your Age Group:", list(age_groups.keys()))

# Initialize session
if "session_initialized" not in st.session_state or st.session_state.age_group != selected_age_group:
    st.session_state.age_group = selected_age_group
    st.session_state.category = None
    st.session_state.question_index = 0
    st.session_state.show_answer = False
    st.session_state.show_steps = False
    st.session_state.session_initialized = True

# Load dataset
group_info = age_groups[selected_age_group]
dataset_path = group_info["dataset"]
zip_path = group_info["zip_file"]
image_folder = group_info["image_folder"]

# Create image folder in temp directory
temp_image_folder = os.path.join(st.session_state.temp_dir, image_folder)
os.makedirs(temp_image_folder, exist_ok=True)

# Extract images if zip exists
try:
    if os.path.exists(zip_path):
        with zipfile.ZipFile(zip_path, "r") as zip_ref:
            zip_ref.extractall(temp_image_folder)
    else:
        st.warning(f"Zip file not found: {zip_path}. Using default images if available.")
except PermissionError as e:
    st.error(f"Permission error accessing files: {e}")
    st.info("Using default images without extraction.")
except Exception as e:
    st.error(f"Error extracting images: {e}")

if not os.path.exists(dataset_path):
    st.error(f"Dataset not found: {dataset_path}")
    st.stop()

try:
    df = pd.read_excel(dataset_path)
    df['category'] = df['category'].astype(str).str.strip()
except Exception as e:
    st.error(f"Error loading dataset: {e}")
    st.stop()

# Category selection
categories = sorted(df['category'].dropna().unique())
selected_category = st.selectbox("📚 Choose a Math Category:", options=categories)

# Update session if category changes
if st.session_state.category != selected_category:
    st.session_state.category = selected_category
    st.session_state.question_index = 0
    st.session_state.show_answer = False
    st.session_state.show_steps = False
    st.rerun()

# Filter questions by selected category only
subset_df = df[df['category'] == selected_category].reset_index(drop=True)

if not subset_df.empty and st.session_state.question_index < len(subset_df):
    question = subset_df.iloc[st.session_state.question_index]
    progress = int((st.session_state.question_index / len(subset_df)) * 100)
    st.progress(progress)

    st.markdown(f"<div class='question-box'>📘 <b>Question {st.session_state.question_index + 1}:</b><br><br>{question['problem']}</div>", unsafe_allow_html=True)

    # Display image with error handling
    if pd.notna(question.get('image')):
        image_name = str(question['image']).strip()
        image_found = False
        
        # Check in temp directory first
        for root, _, files in os.walk(temp_image_folder):
            for file in files:
                if file.lower().startswith(image_name.lower()) or os.path.splitext(file)[0].lower() == image_name.lower():
                    try:
                        img_path = os.path.join(root, file)
                        st.image(Image.open(img_path), use_column_width=True)
                        image_found = True
                        break
                    except Exception as e:
                        st.warning(f"❌ Couldn't load image: {e}")
        
        if not image_found:
            # Fallback to original image folder
            for root, _, files in os.walk(image_folder):
                for file in files:
                    if file.lower().startswith(image_name.lower()) or os.path.splitext(file)[0].lower() == image_name.lower():
                        try:
                            st.image(Image.open(os.path.join(root, file)), use_column_width=True)
                            image_found = True
                            break
                        except Exception as e:
                            st.warning(f"❌ Couldn't load image: {e}")
        
        if not image_found:
            st.warning("❌ Image not found.")

    user_ans = st.text_input("📝 Your Answer:", key=f"ans_{st.session_state.question_index}")

    if st.button("✅ Submit Answer"):
        if str(user_ans).strip().lower() == str(question['answer']).strip().lower():
            st.success("🎉 Correct! Well done!")
            st.balloons()
            time.sleep(2)
            st.session_state.question_index += 1
            st.session_state.show_answer = False
            st.session_state.show_steps = False
            st.rerun()
        else:
            st.error("❌ Try again or view the correct answer below.")
            st.session_state.show_answer = True
            st.session_state.show_steps = False

    if st.session_state.show_answer:
        st.info(f"✅ Correct Answer: **{question['answer']}**")
        if selected_age_group in ["7-9 Age Group", "13-15 Age Group"]:
            if st.button("🔍 Show Steps"):
                st.session_state.show_steps = True
        if st.session_state.show_steps and pd.notna(question.get("steps", None)):
            st.success(f"### 🪄 Steps:\n{question['steps']}")

    if st.button("⏭️ Skip"):
        st.session_state.question_index += 1
        st.session_state.show_answer = False
        st.session_state.show_steps = False
        st.rerun()
elif subset_df.empty:
    st.warning("⚠️ No questions available in this category. Try another one.")
else:
    st.success("🏁 You've completed all questions in this category!")

# -------------------------------------------------------------
# ✅ Enhanced "Ask Any Math Question" Section Using MiniLM Model
# -------------------------------------------------------------

@st.cache_resource
def load_model():
    return SentenceTransformer("all-MiniLM-L6-v2")

model = load_model()

# Use pickle to store and share QA data efficiently
qa_data = {
    "what is the area of a triangle": "Area = (1/2) × base × height",
    "what is the area of a square": "Area = side²",
    "what is the area of a circle": "Area = π × radius²",
    "perimeter of square": "Perimeter = 4 × side",
    "perimeter of rectangle": "Perimeter = 2 × (length + width)",
    "volume of cube": "Volume = side³",
    "volume of sphere": "Volume = (4/3) × π × radius³",
    "what is pi": "π ≈ 3.14159",
    "how to calculate speed": "Speed = Distance ÷ Time",
    "what is the derivative of x squared": "The derivative of x² is 2x."
}

# Store QA data using pickle (for demonstration)
try:
    with open('qa_data.pkl', 'wb') as f:
        pickle.dump(qa_data, f)
    
    # Optional: disassemble and analyze pickle (for debugging)
    with open('qa_data.pkl', 'rb') as f:
        pickletools.dis(f)
except Exception as e:
    st.warning(f"Could not create pickle file: {e}")

qa_questions = list(qa_data.keys())
qa_embeddings = model.encode(qa_questions)

st.markdown("## 💡 Ask Any Math Question")
question_input = st.text_input("Ask a question like 'What is the area of a circle?' or 'Derivative of x^2'")

if question_input:
    user_embedding = model.encode([question_input])
    scores = cosine_similarity(user_embedding, qa_embeddings)[0]
    best_idx = scores.argmax()
    best_score = scores[best_idx]

    if best_score > 0.6:
        st.success(f"🤖 Answer:\n\n**{qa_data[qa_questions[best_idx]]}**")
    else:
        try:
            if "derivative" in question_input.lower():
                expr = question_input.lower().split("of")[-1].strip()
                var = sp.symbols('x')
                derivative = sp.diff(sp.sympify(expr.replace("^", "**")), var)
                st.success(f"🧮 Derivative of {expr} is:\n\n**{derivative}**")
            else:
                result = sp.sympify(question_input.replace("^", "**"))
                simplified = sp.simplify(result)
                st.success(f"✅ Answer: {simplified}")
        except Exception as e:
            st.warning(f"⚠️ Couldn't understand the question. Error: {e}")

# Cleanup function (optional)
def cleanup():
    import shutil
    if 'temp_dir' in st.session_state and os.path.exists(st.session_state.temp_dir):
        shutil.rmtree(st.session_state.temp_dir)

# Register cleanup when app closes
import atexit
atexit.register(cleanup)