Update app.py

app.py

@@ -6,11 +6,8 @@ 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(
@@ -38,265 +35,320 @@ def convert_role_for_gemini(role):
     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'
+    """Determine if a visual aid would be helpful for K-12 concepts"""
+    elementary_keywords = [
+        'add', 'subtract', 'multiply', 'divide', 'times', 'plus', 'minus',
+        'count', 'counting', 'number', 'numbers', 'fraction', 'fractions',
+        'shape', 'shapes', 'circle', 'square', 'triangle', 'rectangle',
+        'money', 'coins', 'dollars', 'cents', 'time', 'clock',
+        'pattern', 'patterns', 'grouping', 'groups', 'tens', 'ones',
+        'place value', 'hundred', 'thousand', 'comparison', 'greater', 'less',
+        'equal', 'equals', 'measurement', 'length', 'height', 'weight'
     ]
     
     combined_text = (user_prompt + " " + ai_response).lower()
-    return any(keyword in combined_text for keyword in visual_keywords)
+    return any(keyword in combined_text for keyword in elementary_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"""
+def create_counting_visual(numbers):
+    """Create visual counting aids using dots/circles"""
     try:
-        user_lower = user_prompt.lower()
+        if not isinstance(numbers, list):
+            numbers = [numbers]
         
-        # 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)
-                ))
+        fig = go.Figure()
+        colors = ['red', 'blue', 'green', 'orange', 'purple', 'yellow']
+        
+        for i, num in enumerate(numbers[:6]):  # Limit to 6 different numbers
+            if num <= 20:  # Only for reasonable counting numbers
+                # Create dots arranged in rows
+                dots_per_row = min(5, num)
+                rows = (num - 1) // dots_per_row + 1
                 
-                # 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")
+                x_positions = []
+                y_positions = []
                 
-                # 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")
+                for dot in range(num):
+                    row = dot // dots_per_row
+                    col = dot % dots_per_row
+                    x_positions.append(col + i * 7)  # Separate groups
+                    y_positions.append(-row + rows - 1)
                 
-                # 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}"
+                    x=x_positions,
+                    y=y_positions,
+                    mode='markers',
+                    marker=dict(
+                        size=20,
+                        color=colors[i],
+                        symbol='circle',
+                        line=dict(width=2, color='black')
+                    ),
+                    name=f'{num} items',
+                    showlegend=True
                 ))
                 
-                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"
+                # Add number label
+                fig.add_annotation(
+                    x=2 + i * 7,
+                    y=-rows - 0.5,
+                    text=str(num),
+                    font=dict(size=24, color=colors[i]),
+                    showlegend=False
                 )
-                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)
+        fig.update_layout(
+            title="Counting Visualization",
+            showlegend=True,
+            xaxis=dict(showgrid=False, showticklabels=False, zeroline=False),
+            yaxis=dict(showgrid=False, showticklabels=False, zeroline=False),
+            height=400,
+            template="simple_white"
+        )
+        
+        return fig
+    except:
+        return None
+
+def create_addition_visual(num1, num2):
+    """Create visual addition using manipulatives"""
+    try:
+        if num1 > 10 or num2 > 10:  # Keep it simple for young learners
+            return None
             
-            # 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²")
+        fig = go.Figure()
+        
+        # First number (red circles)
+        x1 = list(range(num1))
+        y1 = [1] * num1
+        fig.add_trace(go.Scatter(
+            x=x1, y=y1,
+            mode='markers',
+            marker=dict(size=25, color='red', symbol='circle', line=dict(width=2, color='black')),
+            name=f'First group: {num1}',
+            showlegend=True
+        ))
+        
+        # Second number (blue circles)
+        x2 = list(range(num1 + 1, num1 + num2 + 1))
+        y2 = [1] * num2
+        fig.add_trace(go.Scatter(
+            x=x2, y=y2,
+            mode='markers',
+            marker=dict(size=25, color='blue', symbol='circle', line=dict(width=2, color='black')),
+            name=f'Second group: {num2}',
+            showlegend=True
+        ))
+        
+        # Plus sign
+        fig.add_annotation(
+            x=num1 - 0.5,
+            y=1.5,
+            text="+",
+            font=dict(size=30, color='black'),
+            showlegend=False
+        )
+        
+        # Equals sign and result
+        fig.add_annotation(
+            x=num1 + num2 + 0.5,
+            y=1.5,
+            text="=",
+            font=dict(size=30, color='black'),
+            showlegend=False
+        )
+        
+        fig.add_annotation(
+            x=num1 + num2 + 1.5,
+            y=1.5,
+            text=str(num1 + num2),
+            font=dict(size=30, color='green'),
+            showlegend=False
+        )
+        
+        fig.update_layout(
+            title=f"Addition: {num1} + {num2} = {num1 + num2}",
+            showlegend=True,
+            xaxis=dict(showgrid=False, showticklabels=False, zeroline=False, range=[-1, num1 + num2 + 3]),
+            yaxis=dict(showgrid=False, showticklabels=False, zeroline=False, range=[0, 2.5]),
+            height=300,
+            template="simple_white"
+        )
+        
+        return fig
+    except:
+        return None
+
+def create_fraction_visual(numerator, denominator):
+    """Create visual fraction using pie charts or bars"""
+    try:
+        if denominator > 12 or numerator > denominator:  # Keep it simple
+            return None
             
-            # 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³")
+        fig = go.Figure()
+        
+        # Create a circle divided into parts
+        angles = np.linspace(0, 2*np.pi, denominator + 1)
+        
+        for i in range(denominator):
+            # Create each slice
+            theta = np.linspace(angles[i], angles[i+1], 50)
+            r = np.ones_like(theta)
+            x = r * np.cos(theta)
+            y = r * np.sin(theta)
             
-            # 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}")
+            # Add center point
+            x = np.concatenate([[0], x, [0]])
+            y = np.concatenate([[0], y, [0]])
             
-            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")
+            color = 'lightblue' if i < numerator else 'lightgray'
             
-            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
+            fig.add_trace(go.Scatter(
+                x=x, y=y,
+                fill='toself',
+                mode='lines',
+                line=dict(color='black', width=2),
+                fillcolor=color,
+                name=f'Slice {i+1}' if i < numerator else '',
+                showlegend=False
+            ))
         
-        # 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
-                )
+        fig.update_layout(
+            title=f"Fraction: {numerator}/{denominator}",
+            xaxis=dict(showgrid=False, showticklabels=False, zeroline=False, scaleanchor="y", scaleratio=1),
+            yaxis=dict(showgrid=False, showticklabels=False, zeroline=False),
+            height=400,
+            template="simple_white"
+        )
+        
+        return fig
+    except:
+        return None
+
+def create_place_value_visual(number):
+    """Create place value visualization"""
+    try:
+        if number > 9999:  # Keep it reasonable for elementary
+            return None
             
-            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)
+        str_num = str(number).zfill(4)  # Pad with zeros
+        ones = int(str_num[-1])
+        tens = int(str_num[-2])
+        hundreds = int(str_num[-3])
+        thousands = int(str_num[-4])
+        
+        fig = go.Figure()
+        
+        # Create visual blocks for each place value
+        positions = [0, 2, 4, 6]  # x positions for thousands, hundreds, tens, ones
+        values = [thousands, hundreds, tens, ones]
+        labels = ['Thousands', 'Hundreds', 'Tens', 'Ones']
+        colors = ['red', 'blue', 'green', 'orange']
+        
+        for i, (pos, val, label, color) in enumerate(zip(positions, values, labels, colors)):
+            if val > 0:
+                # Create blocks to represent the value
+                blocks_per_row = min(5, val)
+                rows = (val - 1) // blocks_per_row + 1
                 
-                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]
+                x_coords = []
+                y_coords = []
                 
-                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)))
+                for block in range(val):
+                    row = block // blocks_per_row
+                    col = block % blocks_per_row
+                    x_coords.append(pos + col * 0.3)
+                    y_coords.append(row * 0.3)
                 
-                fig.update_layout(title="Rectangle (5×3)", height=500)
+                fig.add_trace(go.Scatter(
+                    x=x_coords,
+                    y=y_coords,
+                    mode='markers',
+                    marker=dict(
+                        size=15,
+                        color=color,
+                        symbol='square',
+                        line=dict(width=1, color='black')
+                    ),
+                    name=f'{label}: {val}',
+                    showlegend=True
+                ))
             
-            fig.update_layout(
-                xaxis_title="x",
-                yaxis_title="y",
-                showlegend=True,
-                template="plotly_white",
-                xaxis=dict(zeroline=True),
-                yaxis=dict(zeroline=True)
+            # Add place value label
+            fig.add_annotation(
+                x=pos + 0.6,
+                y=-0.5,
+                text=label,
+                font=dict(size=12),
+                showlegend=False
             )
-            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
+            # Add digit
+            fig.add_annotation(
+                x=pos + 0.6,
+                y=-0.8,
+                text=str(val),
+                font=dict(size=16, color=colors[i]),
+                showlegend=False
             )
-            return fig
+        
+        fig.update_layout(
+            title=f"Place Value: {number}",
+            showlegend=True,
+            xaxis=dict(showgrid=False, showticklabels=False, zeroline=False, range=[-0.5, 7]),
+            yaxis=dict(showgrid=False, showticklabels=False, zeroline=False),
+            height=400,
+            template="simple_white"
+        )
+        
+        return fig
+    except:
+        return None
+
+def generate_k12_visual(user_prompt, ai_response):
+    """Generate age-appropriate visualizations for K-12 students"""
+    try:
+        user_lower = user_prompt.lower()
+        
+        # COUNTING NUMBERS
+        count_match = re.search(r'count.*?(\d+)', user_lower)
+        if count_match or 'counting' in user_lower:
+            number = int(count_match.group(1)) if count_match else 5
+            return create_counting_visual(number)
+        
+        # SIMPLE ADDITION
+        add_match = re.search(r'(\d+)\s*\+\s*(\d+)', user_prompt)
+        if add_match and 'add' in user_lower or '+' in user_prompt:
+            num1, num2 = int(add_match.group(1)), int(add_match.group(2))
+            if num1 <= 10 and num2 <= 10:  # Keep it simple
+                return create_addition_visual(num1, num2)
+        
+        # FRACTIONS
+        fraction_match = re.search(r'(\d+)/(\d+)', user_prompt)
+        if fraction_match or 'fraction' in user_lower:
+            if fraction_match:
+                num, den = int(fraction_match.group(1)), int(fraction_match.group(2))
+            else:
+                num, den = 1, 2  # Default to 1/2
+            return create_fraction_visual(num, den)
+        
+        # PLACE VALUE
+        if 'place value' in user_lower or 'place' in user_lower:
+            place_match = re.search(r'\b(\d{1,4})\b', user_prompt)
+            if place_match:
+                number = int(place_match.group(1))
+                return create_place_value_visual(number)
+        
+        # NUMBERS (general counting)
+        number_match = re.search(r'\b(\d+)\b', user_prompt)
+        if number_match and any(word in user_lower for word in ['show', 'count', 'number']):
+            number = int(number_match.group(1))
+            if 1 <= number <= 20:
+                return create_counting_visual(number)
         
         return None
         
     except Exception as e:
-        st.error(f"Could not generate Plotly visual: {e}")
+        st.error(f"Could not generate K-12 visual: {e}")
         return None
 
 # --- API KEY & MODEL CONFIGURATION ---
@@ -315,48 +367,43 @@ if api_key:
     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.
+        You are "Math Jegna", an AI math tutor specializing in K-12 mathematics using the Professor B methodology.
+        
+        FOCUS ON ELEMENTARY CONCEPTS:
+        - Basic counting (1-100)
+        - Simple addition and subtraction (single digits to start)
+        - Beginning multiplication (times tables)
+        - Basic fractions (halves, thirds, quarters)
+        - Place value (ones, tens, hundreds)
+        - Shape recognition
+        - Simple word problems
+        - Money and time concepts
+        
+        PROFESSOR B METHODOLOGY - ESSENTIAL PRINCIPLES:
+        1. Present math as a STORY that connects ideas
+        2. Use MENTAL GYMNASTICS - fun games and finger counting
+        3. Build from concrete to abstract naturally
+        4. NO ROTE MEMORIZATION - focus on understanding patterns and connections
+        5. Eliminate math anxiety through simple, truthful explanations
+        6. Use manipulatives and visual aids
+        7. Allow accelerated learning when student shows mastery
+        
+        TEACHING STYLE:
+        - Start with what the child already knows
+        - Build new concepts as natural extensions of previous learning
+        - Use simple, clear language appropriate for the age
+        - Make math enjoyable and reduce tension
+        - Connect everything to real-world experiences
+        - Celebrate understanding, not just correct answers
+        
+        VISUAL AIDS: Mention when visual aids will help, using phrases like:
+        - "Let me show you this with counting dots..."
+        - "I'll create a picture to help you see this..."
+        - "A visual will make this clearer..."
+        
+        Remember: You're helping young minds discover the beauty and logic of mathematics through stories and connections, not through drilling and memorization.
+        
+        You are strictly forbidden from answering non-mathematical questions. If asked non-math questions, respond only with: "I can only answer mathematical questions. Please ask me a question about counting, adding, shapes, or another math topic."
         """
     )
 else:
@@ -383,7 +430,7 @@ if "chats" not in st.session_state:
         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! 🧠📊"}
+                    {"role": "assistant", "content": "Hello! I'm Math Jegna, your friendly math helper! 🧠✨\n\nI love helping kids learn math through fun stories and pictures. Try asking me about:\n- Counting numbers\n- Adding or subtracting\n- Fractions like 1/2\n- Shapes and patterns\n- Or any math question!\n\nWhat would you like to learn about today?"}
                 ]
             }
             st.session_state.active_chat_key = "New Chat"
@@ -414,7 +461,7 @@ def delete_chat(chat_key):
         st.rerun()
 
 # --- SIDEBAR CHAT MANAGEMENT ---
-st.sidebar.title("📚 My Chats")
+st.sidebar.title("📚 My Math Chats")
 st.sidebar.divider()
 
 if st.sidebar.button("➕ New Chat", use_container_width=True):
@@ -423,7 +470,7 @@ if st.sidebar.button("➕ New Chat", use_container_width=True):
         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. ����📊"}
+        {"role": "assistant", "content": "Hi there! Ready for some fun math learning? Ask me about counting, adding, shapes, or anything else! 🚀🔢"}
     ]
     st.session_state.active_chat_key = new_chat_key
     st.rerun()
@@ -459,20 +506,31 @@ for chat_key in list(st.session_state.chats.keys()):
 # --- 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!")
+st.title(f"Math Helper: {st.session_state.active_chat_key} 🧠")
+st.write("🎯 Perfect for young learners! Ask about counting, adding, shapes, fractions, and more!")
+
+# Add some example prompts for young learners
+with st.expander("💡 Try asking me about..."):
+    st.write("""
+    - **Counting**: "Show me how to count to 10"
+    - **Addition**: "What is 3 + 4?"
+    - **Fractions**: "What is 1/2?"
+    - **Place Value**: "What is the place value of 325?"
+    - **Shapes**: "Tell me about triangles"
+    - **Time**: "How do I read a clock?"
+    """)
 
 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():
+if user_prompt := st.chat_input("Ask me a math question!"):
     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... 🤔"):
+        with st.spinner("Math Jegna is thinking... 🤔"):
             try:
                 # Generate text response first
                 chat_session = model.start_chat(history=[
@@ -488,14 +546,14 @@ if user_prompt := st.chat_input():
                 
                 # 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.")
+                    with st.spinner("Creating a helpful picture... 🎨"):
+                        k12_fig = generate_k12_visual(user_prompt, ai_response_text)
+                        if k12_fig:
+                            st.plotly_chart(k12_fig, use_container_width=True)
+                            st.success("✨ Here's a picture to help you understand!")
 
             except Exception as e:
-                error_message = f"Sorry, something went wrong. Math Mentor is taking a break! 🤖\n\n**Error:** {e}"
+                error_message = f"Oops! Something went wrong. Let me try again! 🤖\n\n**Error:** {e}"
                 st.error(error_message)
                 active_chat.append({"role": "assistant", "content": error_message})