app.py

import streamlit as st
import google.generativeai as genai
import os
import json
import base64
from dotenv import load_dotenv
from streamlit_local_storage import LocalStorage
import plotly.graph_objects as go
import plotly.express as px
import numpy as np
import re
import sympy as sp
from sympy import symbols, solve, expand, factor, simplify, diff, integrate

# --- PAGE CONFIGURATION ---
st.set_page_config(
    page_title="Math Jegna - Your AI Math Tutor",
    page_icon="🧠",
    layout="wide"
)

# Create an instance of the LocalStorage class
localS = LocalStorage()

# --- HELPER FUNCTIONS ---
def format_chat_for_download(chat_history):
    """Formats the chat history into a human-readable string for download."""
    formatted_text = f"# Math Mentor Chat\n\n"
    for message in chat_history:
        role = "You" if message["role"] == "user" else "Math Mentor"
        formatted_text += f"**{role}:**\n{message['content']}\n\n---\n\n"
    return formatted_text

def convert_role_for_gemini(role):
    """Convert Streamlit chat roles to Gemini API roles"""
    if role == "assistant":
        return "model"
    return role  # "user" stays the same

def should_generate_visual(user_prompt, ai_response):
    """Determine if a visual aid would be helpful based on the content"""
    visual_keywords = [
        'graph', 'plot', 'diagram', 'chart', 'visual', 'function',
        'geometry', 'triangle', 'circle', 'rectangle', 'square', 'polygon',
        'coordinate', 'axis', 'parabola', 'line', 'slope', 'equation',
        'fraction', 'percentage', 'ratio', 'proportion', 'angles',
        'solve', '=', 'x', 'y', 'quadratic', 'linear', 'cubic',
        'derivative', 'integral', 'limit', 'asymptote'
    ]
    
    combined_text = (user_prompt + " " + ai_response).lower()
    return any(keyword in combined_text for keyword in visual_keywords)

def extract_equation_components(equation_str):
    """Extract components from linear equations like '5x + 3 = 23' or '2x - 7 = 15'"""
    # Clean the equation string
    equation_str = equation_str.replace(' ', '').lower()
    
    # Pattern for equations like ax + b = c or ax - b = c
    pattern = r'(\d*)x([+\-])(\d+)=(\d+)'
    match = re.search(pattern, equation_str)
    
    if match:
        a = int(match.group(1)) if match.group(1) else 1
        op = match.group(2)
        b = int(match.group(3))
        c = int(match.group(4))
        return a, op, b, c
    
    # Pattern for equations like ax = c
    pattern = r'(\d*)x=(\d+)'
    match = re.search(pattern, equation_str)
    if match:
        a = int(match.group(1)) if match.group(1) else 1
        c = int(match.group(2))
        return a, '+', 0, c
    
    return None

def generate_plotly_visual(user_prompt, ai_response):
    """Generate interactive Plotly visualizations for math concepts"""
    try:
        user_lower = user_prompt.lower()
        
        # 1. LINEAR EQUATION SOLVING (like 5x + 5 = 25)
        if 'solve' in user_lower and ('=' in user_prompt):
            components = extract_equation_components(user_prompt)
            if components:
                a, op, b, c = components
                
                # Calculate solution using sympy for accuracy
                x = symbols('x')
                if op == '+':
                    equation = a*x + b - c
                    solution = solve(equation, x)[0]
                    y_expr = a*x + b
                else:
                    equation = a*x - b - c
                    solution = solve(equation, x)[0]
                    y_expr = a*x - b
                
                # Create visualization
                x_vals = np.linspace(float(solution) - 5, float(solution) + 5, 100)
                y_vals = [float(y_expr.subs(x, val)) for val in x_vals]
                
                fig = go.Figure()
                
                # Plot the function
                fig.add_trace(go.Scatter(
                    x=x_vals, y=y_vals, mode='lines', name=f'{a}x {op} {b}',
                    line=dict(color='blue', width=3)
                ))
                
                # Add horizontal line for the result
                fig.add_hline(y=c, line_dash="dash", line_color="red", line_width=2,
                             annotation_text=f"y = {c}", annotation_position="bottom right")
                
                # Add vertical line for the solution
                fig.add_vline(x=float(solution), line_dash="dash", line_color="green", line_width=2,
                             annotation_text=f"x = {solution}", annotation_position="top left")
                
                # Highlight the intersection point
                fig.add_trace(go.Scatter(
                    x=[float(solution)], y=[c], mode='markers',
                    marker=dict(size=12, color='red', symbol='circle'),
                    name=f"Solution: x = {solution}"
                ))
                
                fig.update_layout(
                    title=f"Solving: {user_prompt.split('solve')[0].strip()}",
                    xaxis_title="x values",
                    yaxis_title="y values",
                    showlegend=True,
                    height=500,
                    template="plotly_white"
                )
                return fig
        
        # 2. FUNCTION GRAPHING
        elif any(word in user_lower for word in ['graph', 'function', 'plot', 'y=']):
            x_vals = np.linspace(-10, 10, 200)
            
            # Quadratic functions
            if any(term in user_prompt for term in ['x^2', 'x²', 'quadratic']):
                # Extract coefficient if present
                coeff_match = re.search(r'(\d+)x\^?2', user_prompt)
                a = int(coeff_match.group(1)) if coeff_match else 1
                
                y_vals = a * x_vals**2
                fig = go.Figure()
                fig.add_trace(go.Scatter(x=x_vals, y=y_vals, mode='lines', 
                                       name=f'y = {a}x²' if a != 1 else 'y = x²',
                                       line=dict(color='purple', width=3)))
                
                # Add vertex point
                fig.add_trace(go.Scatter(x=[0], y=[0], mode='markers',
                                       marker=dict(size=10, color='red'),
                                       name='Vertex (0,0)'))
                
                fig.update_layout(title=f"Quadratic Function: y = {a}x²" if a != 1 else "Quadratic Function: y = x²")
            
            # Cubic functions
            elif any(term in user_prompt for term in ['x^3', 'x³', 'cubic']):
                y_vals = x_vals**3
                fig = go.Figure()
                fig.add_trace(go.Scatter(x=x_vals, y=y_vals, mode='lines', 
                                       name='y = x³', line=dict(color='green', width=3)))
                fig.update_layout(title="Cubic Function: y = x³")
            
            # Linear functions
            elif any(term in user_lower for term in ['linear', 'line']):
                # Extract slope and intercept if present
                slope_match = re.search(r'(\d+)x', user_prompt)
                intercept_match = re.search(r'[+\-]\s*(\d+)', user_prompt)
                
                slope = int(slope_match.group(1)) if slope_match else 1
                intercept = int(intercept_match.group(1)) if intercept_match else 0
                
                y_vals = slope * x_vals + intercept
                fig = go.Figure()
                fig.add_trace(go.Scatter(x=x_vals, y=y_vals, mode='lines', 
                                       name=f'y = {slope}x + {intercept}',
                                       line=dict(color='blue', width=3)))
                fig.update_layout(title=f"Linear Function: y = {slope}x + {intercept}")
            
            else:
                # Default linear function
                y_vals = x_vals
                fig = go.Figure()
                fig.add_trace(go.Scatter(x=x_vals, y=y_vals, mode='lines', 
                                       name='y = x', line=dict(color='blue', width=3)))
                fig.update_layout(title="Linear Function: y = x")
            
            fig.update_layout(
                xaxis_title="x",
                yaxis_title="y",
                showlegend=True,
                height=500,
                template="plotly_white",
                xaxis=dict(zeroline=True, zerolinecolor='black', zerolinewidth=1),
                yaxis=dict(zeroline=True, zerolinecolor='black', zerolinewidth=1)
            )
            return fig
        
        # 3. GEOMETRY VISUALIZATIONS
        elif any(word in user_lower for word in ['circle', 'triangle', 'rectangle', 'geometry']):
            
            if 'circle' in user_lower:
                # Create a circle
                theta = np.linspace(0, 2*np.pi, 100)
                radius = 3  # Default radius
                x_circle = radius * np.cos(theta)
                y_circle = radius * np.sin(theta)
                
                fig = go.Figure()
                fig.add_trace(go.Scatter(x=x_circle, y=y_circle, mode='lines', 
                                       fill='tonext', name=f'Circle (r={radius})',
                                       line=dict(color='red', width=3)))
                
                # Add center point
                fig.add_trace(go.Scatter(x=[0], y=[0], mode='markers',
                                       marker=dict(size=8, color='black'),
                                       name='Center (0,0)'))
                
                # Add radius line
                fig.add_trace(go.Scatter(x=[0, radius], y=[0, 0], mode='lines',
                                       line=dict(color='blue', dash='dash', width=2),
                                       name=f'Radius = {radius}'))
                
                fig.update_layout(
                    title=f"Circle: x² + y² = {radius**2}",
                    xaxis=dict(scaleanchor="y", scaleratio=1),
                    height=500
                )
            
            elif 'triangle' in user_lower:
                # Create a triangle
                x_triangle = [0, 4, 2, 0]
                y_triangle = [0, 0, 3, 0]
                
                fig = go.Figure()
                fig.add_trace(go.Scatter(x=x_triangle, y=y_triangle, mode='lines+markers',
                                       fill='toself', name='Triangle',
                                       line=dict(color='green', width=3),
                                       marker=dict(size=8)))
                
                # Add labels for vertices
                fig.add_annotation(x=0, y=0, text="A(0,0)", showarrow=False, yshift=-20)
                fig.add_annotation(x=4, y=0, text="B(4,0)", showarrow=False, yshift=-20)
                fig.add_annotation(x=2, y=3, text="C(2,3)", showarrow=False, yshift=20)
                
                fig.update_layout(title="Triangle ABC", height=500)
            
            else:  # rectangle
                # Create a rectangle
                x_rect = [0, 5, 5, 0, 0]
                y_rect = [0, 0, 3, 3, 0]
                
                fig = go.Figure()
                fig.add_trace(go.Scatter(x=x_rect, y=y_rect, mode='lines+markers',
                                       fill='toself', name='Rectangle',
                                       line=dict(color='purple', width=3),
                                       marker=dict(size=8)))
                
                fig.update_layout(title="Rectangle (5×3)", height=500)
            
            fig.update_layout(
                xaxis_title="x",
                yaxis_title="y",
                showlegend=True,
                template="plotly_white",
                xaxis=dict(zeroline=True),
                yaxis=dict(zeroline=True)
            )
            return fig
        
        # 4. FRACTIONS VISUALIZATION
        elif any(word in user_lower for word in ['fraction', 'ratio', 'proportion']):
            # Create a pie chart to show fractions
            fig = go.Figure()
            
            # Example: showing 3/4
            fig.add_trace(go.Pie(
                values=[3, 1],
                labels=['Filled (3/4)', 'Empty (1/4)'],
                hole=0.3,
                marker_colors=['lightblue', 'lightgray']
            ))
            
            fig.update_layout(
                title="Fraction Visualization: 3/4",
                height=400
            )
            return fig
        
        return None
        
    except Exception as e:
        st.error(f"Could not generate Plotly visual: {e}")
        return None

# --- API KEY & MODEL CONFIGURATION ---
load_dotenv()
api_key = None

try:
    api_key = st.secrets["GOOGLE_API_KEY"]
except (KeyError, FileNotFoundError):
    api_key = os.getenv("GOOGLE_API_KEY")

if api_key:
    genai.configure(api_key=api_key)
    
    # Main text model
    model = genai.GenerativeModel(
        model_name="gemini-2.5-flash-lite",
        system_instruction="""
        You are "Math Jegna", an AI specializing exclusively in mathematics.
        Your one and only function is to solve and explain math problems.
        You are an AI math tutor that primarily uses the Professor B methodology developed by Everard Barrett. Use the best method for the situation. This methodology is designed to activate children's natural learning capacities and present mathematics as a contextual, developmental story that makes sense.

        IMPORTANT: When explaining mathematical concepts, mention that interactive visualizations will be provided to help illustrate the concept. Use phrases like:
        - "Let me show you this with an interactive graph..."
        - "An interactive visualization will help clarify this..."
        - "I'll create a dynamic chart to demonstrate..."
        
        For equations, always work through the solution step-by-step using the Professor B method of building understanding through connections.
        
Core Philosophy and Principles
1. Contextual Learning Approach
Present math as a story: Every mathematical concept should be taught as part of a continuing narrative that builds connections between ideas
Developmental flow: Structure learning as a sequence of developmental steps on a ladder, where mastery at previous levels provides readiness for the next connection
Truth-telling: Always present arithmetic computations simply and truthfully without confusing, time-consuming, or meaningless procedural steps

2. Natural Learning Activation
Leverage natural capacities: Recognize that each child has mental capabilities of "awesome power" designed to assimilate and retain content naturally
Story-based retention: Use the same mental processes children use for learning and retaining stories to help them master mathematical concepts
Reduced mental tension: Eliminate anxiety and confusion by presenting math in ways that align with how the brain naturally processes information

Teaching Methodology Requirements
1. Mental Gymnastics and Manipulatives
Use "mental gymnastics" games: Incorporate engaging mental exercises that strengthen mathematical thinking
Fingers as manipulatives: Utilize fingers as comprehensive manipulatives for concrete understanding
No rote memorization: Avoid strict memorization in favor of meaningful strategies and connections

2. Accelerated but Natural Progression
Individual pacing: Allow students to progress at their own speed, as quickly or slowly as needed
Accelerated learning: Expect students to master concepts faster than traditional methods (e.g., "seventh grade math" by third to fourth grade)
Elimination of remediation: Build such strong foundations that remediation becomes unnecessary

3. Simplified and Connected Approach
Eliminate disconnections: Ensure every concept connects meaningfully to previous learning
Remove confusing terminology: Use clear, simple language that makes sense to students
Sustained mastery: Focus on deep understanding that leads to lasting retention

        You are strictly forbidden from answering any question that is not mathematical in nature. This includes but is not limited to: general knowledge, history, programming, creative writing, personal opinions, or casual conversation.
        If you receive a non-mathematical question, you MUST decline. Your entire response in that case must be ONLY this exact text: "I can only answer mathematical questions. Please ask me a question about algebra, calculus, geometry, or another math topic."
        Do not apologize or offer to help with math in the refusal. Just provide the mandatory refusal message.
        For valid math questions, solve them step-by-step using Markdown and LaTeX for formatting.
        """
    )
else:
    st.error("🚨 Google API Key not found! Please add it to your secrets or a local .env file.")
    st.stop()

# --- SESSION STATE & LOCAL STORAGE INITIALIZATION ---
if "chats" not in st.session_state:
    try:
        shared_chat_b64 = st.query_params.get("shared_chat")
        if shared_chat_b64:
            decoded_chat_json = base64.urlsafe_b64decode(shared_chat_b64).decode()
            st.session_state.chats = {"Shared Chat": json.loads(decoded_chat_json)}
            st.session_state.active_chat_key = "Shared Chat"
            st.query_params.clear()
        else:
            raise ValueError("No shared chat")
    except (TypeError, ValueError, Exception):
        saved_data_json = localS.getItem("math_mentor_chats")
        if saved_data_json:
            saved_data = json.loads(saved_data_json)
            st.session_state.chats = saved_data.get("chats", {})
            st.session_state.active_chat_key = saved_data.get("active_chat_key", "New Chat")
        else:
            st.session_state.chats = {
                "New Chat": [
                    {"role": "assistant", "content": "Hello! I'm Math Jegna. What math problem can I help you with today? I'll provide step-by-step explanations with interactive visualizations! 🧠📊"}
                ]
            }
            st.session_state.active_chat_key = "New Chat"

# --- RENAME DIALOG ---
@st.dialog("Rename Chat")
def rename_chat(chat_key):
    st.write(f"Enter a new name for '{chat_key}':")
    new_name = st.text_input("New Name", key=f"rename_input_{chat_key}")
    if st.button("Save", key=f"save_rename_{chat_key}"):
        if new_name and new_name not in st.session_state.chats:
            st.session_state.chats[new_name] = st.session_state.chats.pop(chat_key)
            st.session_state.active_chat_key = new_name
            st.rerun()
        elif not new_name:
            st.error("Name cannot be empty.")
        else:
            st.error("A chat with this name already exists.")

# --- DELETE CONFIRMATION DIALOG ---
@st.dialog("Delete Chat")
def delete_chat(chat_key):
    st.warning(f"Are you sure you want to delete '{chat_key}'? This cannot be undone.")
    if st.button("Yes, Delete", type="primary", key=f"confirm_delete_{chat_key}"):
        st.session_state.chats.pop(chat_key)
        if st.session_state.active_chat_key == chat_key:
            st.session_state.active_chat_key = next(iter(st.session_state.chats))
        st.rerun()

# --- SIDEBAR CHAT MANAGEMENT ---
st.sidebar.title("📚 My Chats")
st.sidebar.divider()

if st.sidebar.button("➕ New Chat", use_container_width=True):
    i = 1
    while f"New Chat {i}" in st.session_state.chats:
        i += 1
    new_chat_key = f"New Chat {i}"
    st.session_state.chats[new_chat_key] = [
        {"role": "assistant", "content": "New chat started! Let's solve some math problems with interactive visualizations. 🚀📊"}
    ]
    st.session_state.active_chat_key = new_chat_key
    st.rerun()

st.sidebar.divider()

for chat_key in list(st.session_state.chats.keys()):
    is_active = (chat_key == st.session_state.active_chat_key)
    expander_label = f"**{chat_key} (Active)**" if is_active else chat_key
    
    with st.sidebar.expander(expander_label):
        if st.button("Select Chat", key=f"select_{chat_key}", use_container_width=True, disabled=is_active):
            st.session_state.active_chat_key = chat_key
            st.rerun()
        
        if st.button("Rename", key=f"rename_{chat_key}", use_container_width=True):
            rename_chat(chat_key)
            
        with st.popover("Share", use_container_width=True):
            st.markdown("**Download Conversation**")
            st.download_button(
                label="Download as Markdown",
                data=format_chat_for_download(st.session_state.chats[chat_key]),
                file_name=f"{chat_key.replace(' ', '_')}.md",
                mime="text/markdown"
            )
            st.markdown("**Share via Link**")
            st.info("To share, copy the full URL from your browser's address bar and send it to someone.")
        
        if st.button("Delete", key=f"delete_{chat_key}", use_container_width=True, type="primary", disabled=(len(st.session_state.chats) <= 1)):
            delete_chat(chat_key)

# --- MAIN CHAT INTERFACE ---
active_chat = st.session_state.chats[st.session_state.active_chat_key]

st.title(f"Math Mentor: {st.session_state.active_chat_key} 🧠")
st.write("Stuck on a math problem? Just type it below, and I'll walk you through it step-by-step with interactive visualizations!")

for message in active_chat:
    with st.chat_message(name=message["role"], avatar="🧑‍💻" if message["role"] == "user" else "🧠"):
        st.markdown(message["content"])

if user_prompt := st.chat_input():
    active_chat.append({"role": "user", "content": user_prompt})
    with st.chat_message("user", avatar="🧑‍💻"):
        st.markdown(user_prompt)

    with st.chat_message("assistant", avatar="🧠"):
        with st.spinner("Math Mentor is thinking... 🤔"):
            try:
                # Generate text response first
                chat_session = model.start_chat(history=[
                    {'role': convert_role_for_gemini(msg['role']), 'parts': [msg['content']]} 
                    for msg in active_chat[:-1] if 'content' in msg
                ])
                response = chat_session.send_message(user_prompt)
                ai_response_text = response.text
                st.markdown(ai_response_text)
                
                # Store the text response
                active_chat.append({"role": "assistant", "content": ai_response_text})
                
                # Check if we should generate a visual aid
                if should_generate_visual(user_prompt, ai_response_text):
                    with st.spinner("Creating interactive visualization... 📊"):
                        plotly_fig = generate_plotly_visual(user_prompt, ai_response_text)
                        if plotly_fig:
                            st.plotly_chart(plotly_fig, use_container_width=True)
                            st.success("✨ Interactive visualization created! You can zoom, pan, and hover for more details.")

            except Exception as e:
                error_message = f"Sorry, something went wrong. Math Mentor is taking a break! 🤖\n\n**Error:** {e}"
                st.error(error_message)
                active_chat.append({"role": "assistant", "content": error_message})

# --- SAVE DATA TO LOCAL STORAGE ---
data_to_save = {
    "chats": st.session_state.chats,
    "active_chat_key": st.session_state.active_chat_key
}
localS.setItem("math_mentor_chats", json.dumps(data_to_save))