app.py

import streamlit as st
import time
from datetime import datetime, timedelta
from pathlib import Path
import threading
import logging
from typing import List, Dict, Any

# Import our modules
from config.settings import settings
from capture.camera import Camera
from vision.detector import ObjectDetector
from memory.database import MemoryDatabase
from api.openai_client import OpenAIClient

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Initialize components
@st.cache_resource
def initialize_components():
    """Initialize all system components."""
    try:
        # Validate configuration
        settings.validate_config()
        
        # Initialize components
        camera = Camera()
        detector = ObjectDetector()
        database = MemoryDatabase()
        openai_client = OpenAIClient()
        
        # Test OpenAI connection
        if not openai_client.validate_api_connection():
            st.error("⚠️ OpenAI API connection failed. Please check your API key.")
        
        return camera, detector, database, openai_client
        
    except Exception as e:
        st.error(f"❌ Failed to initialize components: {e}")
        return None, None, None, None

def capture_and_process_image(camera: Camera, detector: ObjectDetector, 
                            database: MemoryDatabase, openai_client: OpenAIClient):
    """Capture and process a single image."""
    try:
        # Capture image
        result = camera.capture_and_save()
        if result is None:
            logger.error("Failed to capture image")
            return
        
        image_path, timestamp = result
        
        # Detect objects
        detection_results = detector.detect_objects(image_path)
        if not detection_results:
            logger.error("Failed to detect objects")
            return
        
        # Enhance description with ChatGPT
        objects = detection_results.get('unique_objects', [])
        basic_description = detection_results.get('scene_description', '')
        enhanced_description = openai_client.enhance_scene_description(
            objects, basic_description
        )
        
        # Analyze context
        context_info = detector.analyze_scene_context(
            detection_results.get('detections', [])
        )
        
        # Store in database
        memory_id = database.add_memory(
            image_path=str(image_path),
            timestamp=timestamp,
            scene_description=enhanced_description,
            objects_detected=objects,
            context_info=context_info
        )
        
        logger.info(f"Memory processed and stored with ID: {memory_id}")
        
    except Exception as e:
        logger.error(f"Error in capture and process: {e}")

def search_memories(query: str, database: MemoryDatabase, 
                   openai_client: OpenAIClient) -> List[Dict[str, Any]]:
    """Search memories using natural language query."""
    try:
        # Analyze query with ChatGPT
        query_analysis = openai_client.analyze_query(query)
        
        # Extract search parameters
        objects = query_analysis.get('objects', [])
        time_reference = query_analysis.get('time_reference')
        
        # Convert time reference to date range
        start_date = None
        end_date = None
        
        if time_reference:
            if 'today' in time_reference.lower():
                start_date = datetime.now().replace(hour=0, minute=0, second=0, microsecond=0)
                end_date = datetime.now()
            elif 'yesterday' in time_reference.lower():
                yesterday = datetime.now() - timedelta(days=1)
                start_date = yesterday.replace(hour=0, minute=0, second=0, microsecond=0)
                end_date = yesterday.replace(hour=23, minute=59, second=59)
            elif 'last week' in time_reference.lower():
                start_date = datetime.now() - timedelta(days=7)
                end_date = datetime.now()
        
        # Search database
        memories = database.search_memories(
            query=query,
            start_date=start_date,
            end_date=end_date,
            objects=objects if objects else None,
            limit=20
        )
        
        # Log query
        database.add_query_log(query, len(memories))
        
        return memories
        
    except Exception as e:
        logger.error(f"Error searching memories: {e}")
        return []

def main():
    st.set_page_config(
        page_title="Vision-based Memory Assistant",
        page_icon="🧠",
        layout="wide"
    )
    
    st.title("🧠 Vision-based Personal Memory Assistant")
    st.markdown("A prototype AI system that helps recall visual memories using images captured throughout the day.")
    
    # Initialize components
    camera, detector, database, openai_client = initialize_components()
    
    if not all([camera, detector, database, openai_client]):
        st.error("❌ System initialization failed. Please check your configuration.")
        return
    
    # Sidebar for controls
    st.sidebar.header("🎛️ Controls")
    
    # Manual capture button
    if st.sidebar.button("📸 Capture Image Now"):
        with st.spinner("Capturing and processing image..."):
            capture_and_process_image(camera, detector, database, openai_client)
        st.success("✅ Image captured and processed!")
        st.rerun()
    
    # Auto-capture toggle
    auto_capture = st.sidebar.checkbox("🔄 Enable Auto-Capture", value=False)
    
    if auto_capture:
        st.sidebar.info("Auto-capture is enabled. Images will be captured every 5 minutes during active hours.")
    
    # Main content area
    tab1, tab2, tab3 = st.tabs(["🔍 Search Memories", "📊 Statistics", "📸 Recent Captures"])
    
    with tab1:
        st.header("🔍 Search Your Visual Memories")
        
        # Search interface
        col1, col2 = st.columns([3, 1])
        
        with col1:
            query = st.text_input(
                "Ask about your memories:",
                placeholder="When did I last see my keys? Show me when I was working at my desk..."
            )
        
        with col2:
            search_button = st.button("🔍 Search", type="primary")
        
        if search_button and query:
            with st.spinner("Searching memories..."):
                memories = search_memories(query, database, openai_client)
                
                if memories:
                    # Generate AI response
                    ai_response = openai_client.generate_memory_response(query, memories)
                    st.success(f"🤖 {ai_response}")
                    
                    # Display results
                    st.subheader(f"📋 Found {len(memories)} memories:")
                    
                    for i, memory in enumerate(memories):
                        with st.expander(f"📅 {memory['timestamp'].strftime('%B %d at %I:%M %p')}"):
                            col1, col2 = st.columns([1, 2])
                            
                            with col1:
                                # Display image if it exists
                                image_path = Path(memory['image_path'])
                                if image_path.exists():
                                    st.image(str(image_path), caption="Captured Image")
                                else:
                                    st.warning("Image file not found")
                            
                            with col2:
                                st.write(f"**Scene:** {memory['scene_description']}")
                                st.write(f"**Objects:** {', '.join(memory['objects_detected'])}")
                                
                                if memory['context_info']:
                                    context = memory['context_info']
                                    st.write(f"**Location:** {context.get('location_type', 'Unknown')}")
                                    st.write(f"**Activity:** {context.get('activity_type', 'Unknown')}")
                else:
                    st.warning("No memories found matching your query.")
                    
                    # Suggest alternative searches
                    suggestions = openai_client.suggest_search_terms(query)
                    if suggestions:
                        st.info("💡 Try these alternative searches:")
                        for suggestion in suggestions:
                            if st.button(suggestion, key=f"suggest_{suggestion}"):
                                st.session_state.query = suggestion
                                st.rerun()
        
        # Quick search examples
        st.subheader("💡 Example Queries")
        examples = [
            "When did I last see my laptop?",
            "Show me when I was in the kitchen",
            "Find memories from today",
            "When was I working at my desk?",
            "Show me outdoor scenes"
        ]
        
        cols = st.columns(len(examples))
        for i, example in enumerate(examples):
            with cols[i]:
                if st.button(example, key=f"example_{i}"):
                    st.session_state.query = example
                    st.rerun()
    
    with tab2:
        st.header("📊 Memory Statistics")
        
        stats = database.get_statistics()
        
        col1, col2, col3 = st.columns(3)
        
        with col1:
            st.metric("Total Memories", stats.get('total_memories', 0))
        
        with col2:
            st.metric("Memories Today", stats.get('memories_today', 0))
        
        with col3:
            st.metric("Active Hours", f"{settings.CAPTURE_ACTIVE_HOURS_START}:00 - {settings.CAPTURE_ACTIVE_HOURS_END}:00")
        
        # Most common objects
        if stats.get('most_common_objects'):
            st.subheader("🏷️ Most Common Objects")
            for obj, count in stats['most_common_objects']:
                st.write(f"• {obj}: {count} times")
        
        # Recent query history
        st.subheader("🔍 Recent Searches")
        query_history = database.get_query_history(10)
        
        for query_record in query_history:
            st.write(f"• **{query_record['query_text']}** ({query_record['results_count']} results) - {query_record['timestamp'].strftime('%H:%M')}")
    
    with tab3:
        st.header("📸 Recent Captures")
        
        recent_memories = database.get_recent_memories(10)
        
        if recent_memories:
            for memory in recent_memories:
                with st.expander(f"📅 {memory['timestamp'].strftime('%B %d at %I:%M %p')}"):
                    col1, col2 = st.columns([1, 2])
                    
                    with col1:
                        image_path = Path(memory['image_path'])
                        if image_path.exists():
                            st.image(str(image_path), caption="Captured Image")
                        else:
                            st.warning("Image file not found")
                    
                    with col2:
                        st.write(f"**Scene:** {memory['scene_description']}")
                        st.write(f"**Objects:** {', '.join(memory['objects_detected'])}")
        else:
            st.info("No memories captured yet. Try capturing an image manually or enable auto-capture.")
    
    # Footer
    st.markdown("---")
    st.markdown("""
    <div style='text-align: center; color: #666;'>
        <p>🧠 Vision-based Personal Memory Assistant | Prototype</p>
        <p>Built with PyTorch, YOLO, OpenAI, and Streamlit</p>
    </div>
    """, unsafe_allow_html=True)

if __name__ == "__main__":
    main()