streamlit_app.py

import streamlit as st
from PIL import Image
import os
from dotenv import load_dotenv

# Compatibility fix for Pillow 10.0.0+ where ANTIALIAS was removed
if not hasattr(Image, 'ANTIALIAS'):
    Image.ANTIALIAS = Image.LANCZOS

from utils.ocr import OCRProcessor
from utils.audio import AudioProcessor
from utils.memory import MemorySystem
from utils.hitl import HITLSystem
from agents.parser import ParserAgent
from agents.router import RouterAgent
from agents.solver import SolverAgent
from agents.verifier import VerifierAgent
from agents.explainer import ExplainerAgent
from rag.retriever import Retriever

load_dotenv()

st.set_page_config(page_title="Math Mentor", page_icon="📐", layout="wide")

if 'memory' not in st.session_state:
    st.session_state.memory = MemorySystem()
if 'hitl' not in st.session_state:
    st.session_state.hitl = HITLSystem()
if 'ocr' not in st.session_state:
    st.session_state.ocr = OCRProcessor()
if 'audio' not in st.session_state:
    st.session_state.audio = AudioProcessor()
if 'parser' not in st.session_state:
    st.session_state.parser = ParserAgent()
if 'router' not in st.session_state:
    st.session_state.router = RouterAgent()
if 'solver' not in st.session_state:
    st.session_state.solver = SolverAgent()
if 'verifier' not in st.session_state:
    st.session_state.verifier = VerifierAgent()
if 'explainer' not in st.session_state:
    st.session_state.explainer = ExplainerAgent()
if 'retriever' not in st.session_state:
    st.session_state.retriever = Retriever()

st.title("📐 Math Mentor - AI Problem Solver")
st.markdown("Upload an image, record audio, or type your math problem")

col1, col2 = st.columns([2, 1])

with col1:
    input_mode = st.radio("Input Mode", ["Text", "Image", "Audio"], horizontal=True)
    
    extracted_text = ""
    needs_review = False
    confidence = 1.0
    ocr_confidence = 1.0
    audio_confidence = 1.0
    
    if input_mode == "Text":
        extracted_text = st.text_area("Extracted Text (edit if needed):", value=extracted_text, height=150, key="edited_text")
        
    elif input_mode == "Image":
        uploaded_file = st.file_uploader("Upload image", type=['png', 'jpg', 'jpeg'])
        if uploaded_file:
            image = Image.open(uploaded_file)
            st.image(image, caption="Uploaded Image", use_column_width=True)
            
            with st.spinner("Extracting text from image..."):
                result = st.session_state.ocr.extract_text(image)
                extracted_text = result['text']
                ocr_confidence = result['confidence']
                needs_review = result['needs_review']
            
            col_conf1, col_conf2 = st.columns(2)
            with col_conf1:
                st.metric("OCR Confidence", f"{ocr_confidence:.2%}")
            with col_conf2:
                if needs_review:
                    st.error("⚠️ Low Confidence")
                else:
                    st.success("✅ High Confidence")
            
            extracted_text = st.text_area("Extracted Text (edit if needed):", value=extracted_text, height=150)
            
    elif input_mode == "Audio":
        audio_file = st.file_uploader("Upload audio file", type=['wav', 'mp3', 'm4a'])
        if audio_file:
            st.audio(audio_file)
            
            with st.spinner("Transcribing audio..."):
                result = st.session_state.audio.transcribe(audio_file)
                extracted_text = result['text']
                audio_confidence = result['confidence']
                needs_review = result['needs_review']
            
            col_conf1, col_conf2 = st.columns(2)
            with col_conf1:
                st.metric("Transcription Confidence", f"{audio_confidence:.2%}")
            with col_conf2:
                if needs_review:
                    st.error("⚠️ Low Confidence")
                else:
                    st.success("✅ High Confidence")
            
            extracted_text = st.text_area("Transcription (edit if needed):", value=extracted_text, height=150)
    
    col_btn1, col_btn2 = st.columns(2)
    with col_btn1:
        solve_button = st.button("🚀 Solve Problem", type="primary", disabled=not extracted_text, use_container_width=True)
    with col_btn2:
        recheck_button = st.button("🔍 Request Re-check", disabled=not extracted_text, use_container_width=True)

with col2:
    st.subheader("🔄 Agent Trace")
    trace_container = st.container()

if solve_button and extracted_text:
    final_text = st.session_state.get("edited_text", extracted_text)
    
    # Then use final_text for parsing
    parsed = st.session_state.parser.parse(final_text, input_mode.lower())
    extracted_text = st.session_state.memory.apply_learned_corrections(extracted_text, input_mode.lower())
    
    with trace_container:
        trace = []
        
        st.write("🔍 **Parser Agent**: Analyzing problem...")
        parsed = st.session_state.parser.parse(extracted_text, input_mode.lower())
        trace.append({"agent": "Parser", "output": parsed})
        with st.expander("Parser Output", expanded=False):
            st.json(parsed)
        
        if input_mode == "Image" and extracted_text != result['text']:
            ocr_confidence = 1.0
        if input_mode == "Audio" and extracted_text != result['text']:
            audio_confidence = 1.0

        hitl_check = st.session_state.hitl.should_trigger_hitl(
            ocr_confidence=ocr_confidence,
            audio_confidence=audio_confidence,
            parser_needs_clarification=parsed.get('needs_clarification', False),
            explicit_request=False
        )
        
        if hitl_check['should_trigger']:
            st.error(st.session_state.hitl.get_hitl_instructions(hitl_check))
            st.session_state.hitl_triggered = True
            st.stop()
        
        st.write("🧭 **Router Agent**: Determining strategy...")
        routing = st.session_state.router.route(parsed)
        trace.append({"agent": "Router", "output": routing})
        with st.expander("Router Output", expanded=False):
            st.json(routing)
        
        if routing.get('requires_hitl'):
            st.error(f"❗ HITL Required: {routing.get('reason')}")
            st.info("Please clarify your problem or edit the extracted text.")
            st.stop()
        
        st.write("🔎 **Retriever**: Fetching relevant context...")
        context = st.session_state.retriever.retrieve_context(parsed)
        trace.append({"agent": "Retriever", "sources": len(context['knowledge_base'])})
        st.write(f"📚 Retrieved {len(context['knowledge_base'])} knowledge chunks + {len(context['similar_problems'])} similar problems")
        
        st.write("💡 **Solver Agent**: Solving problem...")
        solution = st.session_state.solver.solve(parsed, context, routing['strategy'])
        trace.append({"agent": "Solver", "steps": len(solution['steps'])})
        if solution.get('calculations_performed', 0) > 0:
            st.write(f"🧮 Performed {solution['calculations_performed']} calculations")
        
        st.write("✅ **Verifier Agent**: Checking solution...")
        verification = st.session_state.verifier.verify(parsed, solution)
        trace.append({"agent": "Verifier", "output": verification})
        
        verifier_hitl = st.session_state.hitl.should_trigger_hitl(
            verifier_confidence=verification.get('confidence', 1.0)
        )
        
        if verifier_hitl['should_trigger']:
            st.warning("⚠️ Verifier has concerns. Solution generated but needs review.")
        
        with st.expander("Verifier Output", expanded=False):
            st.json(verification)
        
        st.write("📚 **Explainer Agent**: Creating explanation...")
        explanation = st.session_state.explainer.explain(parsed, solution, verification)
        
        st.session_state.current_solution = {
            'input_mode': input_mode,
            'original_text': extracted_text,
            'parsed': parsed,
            'routing': routing,
            'solution': solution,
            'verification': verification,
            'explanation': explanation,
            'context': context,
            'trace': trace,
            'hitl_data': verifier_hitl
        }

if recheck_button and extracted_text:
    hitl_explicit = st.session_state.hitl.should_trigger_hitl(explicit_request=True)
    st.warning(st.session_state.hitl.get_hitl_instructions(hitl_explicit))

if 'current_solution' in st.session_state:
    st.markdown("---")
    
    tab1, tab2, tab3, tab4 = st.tabs(["📖 Explanation", "📊 Retrieved Context", "🔍 Solution Details", "📈 Learning Insights"])
    
    with tab1:
        conf = st.session_state.current_solution['verification']['confidence']
        
        col_metric1, col_metric2, col_metric3 = st.columns(3)
        with col_metric1:
            st.metric("Solution Confidence", f"{conf:.2%}")
        with col_metric2:
            is_correct = st.session_state.current_solution['verification']['is_correct']
            st.metric("Verified", "✅ Yes" if is_correct else "⚠️ Review Needed")
        with col_metric3:
            calc_count = st.session_state.current_solution['solution'].get('calculations_performed', 0)
            st.metric("Calculations", calc_count)
        
        if conf < 0.7:
            st.warning("⚠️ Low confidence solution. Please verify carefully.")
        
        st.markdown("### Step-by-Step Explanation")
        st.markdown(st.session_state.current_solution['explanation']['explanation'])
    
    with tab2:
        st.subheader("📚 Knowledge Base Sources")
        for i, item in enumerate(st.session_state.current_solution['context']['knowledge_base'], 1):
            with st.expander(f"Source {i}: {item['metadata']['topic']}"):
                st.write(item['content'])
        
        if st.session_state.current_solution['context']['similar_problems']:
            st.subheader("🔄 Similar Problems from Memory (Self-Learning)")
            for i, prob in enumerate(st.session_state.current_solution['context']['similar_problems'], 1):
                with st.expander(f"Similar Problem {i} (Similarity: {prob.get('similarity', 0):.2%})"):
                    st.write("**Problem:**", prob.get('parsed_question', {}).get('problem_text', ''))
                    if prob.get('user_feedback') == 'correct':
                        st.success("✅ This was a correct solution")
                        if 'solution' in prob:
                            st.write("**Previous Solution:**")
                            st.write(prob['solution'][:300] + "...")
    
    with tab3:
        st.subheader("📝 Full Solution")
        st.write(st.session_state.current_solution['solution']['solution'])
        
        st.subheader("🔍 Verification Results")
        ver = st.session_state.current_solution['verification']
        st.json(ver)
        
        if ver.get('issues'):
            st.subheader("⚠️ Issues Found")
            for issue in ver['issues']:
                st.warning(issue)
    
    with tab4:
        insights = st.session_state.memory.get_learning_insights()
        
        col_i1, col_i2, col_i3 = st.columns(3)
        with col_i1:
            st.metric("Total Problems Solved", insights['total_problems'])
        with col_i2:
            st.metric("Overall Accuracy", f"{insights['accuracy']:.1f}%")
        with col_i3:
            st.metric("Best Strategy", insights['most_successful_strategy'] or "N/A")
        
        if insights['topics_distribution']:
            st.subheader("📊 Topics Distribution")
            st.bar_chart(insights['topics_distribution'])
        
        if insights['common_error_topics']:
            st.subheader("⚠️ Topics Needing Improvement")
            for topic in insights['common_error_topics']:
                st.write(f"- {topic}")
    
    st.markdown("---")
    st.subheader("📝 Provide Feedback (Helps System Learn)")
    
    col_fb1, col_fb2, col_fb3 = st.columns(3)
    
    with col_fb1:
        if st.button("✅ Correct Solution", use_container_width=True):
            st.session_state.memory.store({
                'input_type': st.session_state.current_solution['input_mode'],
                'original_text': st.session_state.current_solution['original_text'],
                'parsed_question': st.session_state.current_solution['parsed'],
                'routing': st.session_state.current_solution['routing'],
                'solution': st.session_state.current_solution['solution']['solution'],
                'verification': st.session_state.current_solution['verification'],
                'user_feedback': 'correct',
                'context_used': st.session_state.current_solution['context']
            })
            st.success("✅ Feedback saved! This solution will help improve future responses.")
            st.balloons()
    
    with col_fb2:
        if st.button("❌ Incorrect Solution", use_container_width=True):
            st.session_state.show_feedback_form = True
    
    with col_fb3:
        if st.button("🔄 Try Again", use_container_width=True):
            del st.session_state.current_solution
            st.rerun()
    
    if st.session_state.get('show_feedback_form'):
        st.markdown("---")
        feedback_comment = st.text_area("What was wrong? Your feedback helps the system learn:", placeholder="e.g., Wrong formula used, calculation error, missed a constraint...")
        
        col_submit, col_cancel = st.columns(2)
        with col_submit:
            if st.button("Submit Feedback", type="primary", use_container_width=True):
                st.session_state.memory.store({
                    'input_type': st.session_state.current_solution['input_mode'],
                    'original_text': st.session_state.current_solution['original_text'],
                    'parsed_question': st.session_state.current_solution['parsed'],
                    'routing': st.session_state.current_solution['routing'],
                    'solution': st.session_state.current_solution['solution']['solution'],
                    'verification': st.session_state.current_solution['verification'],
                    'user_feedback': 'incorrect',
                    'user_comment': feedback_comment,
                    'context_used': st.session_state.current_solution['context']
                })
                st.success("✅ Thank you! This feedback will help the system learn and improve.")
                st.session_state.show_feedback_form = False
                st.rerun()
        
        with col_cancel:
            if st.button("Cancel", use_container_width=True):
                st.session_state.show_feedback_form = False
                st.rerun()

st.sidebar.title("📊 System Statistics")
insights = st.session_state.memory.get_learning_insights()

st.sidebar.metric("Problems Solved", insights['total_problems'])
if insights['total_problems'] > 0:
    st.sidebar.metric("Success Rate", f"{insights['accuracy']:.1f}%")
    
    if insights['most_successful_strategy']:
        st.sidebar.metric("Best Strategy", insights['most_successful_strategy'])
    
    st.sidebar.markdown("### 📚 Topics Learned")
    for topic, count in insights['topics_distribution'].items():
        st.sidebar.write(f"- {topic}: {count}")

st.sidebar.markdown("---")
st.sidebar.info("""
**How This System Learns:**
- Stores all solved problems
- Retrieves similar past solutions
- Learns from your feedback
- Improves OCR/audio corrections
- Identifies successful strategies
""")