wardbuddy/learning_interface.py

"""Gradio interface"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/02_learning_interface.ipynb.

# %% auto 0
__all__ = ['logger', 'create_dashboard_css', 'LearningInterface', 'launch_learning_interface']

# %% ../nbs/02_learning_interface.ipynb 4
from typing import Dict, List, Optional, Tuple, Any
import gradio as gr
from pathlib import Path
import asyncio
from datetime import datetime
import pandas as pd
from .clinical_tutor import ClinicalTutor
from .learning_context import setup_logger

logger = setup_logger(__name__)

# %% ../nbs/02_learning_interface.ipynb 6
def create_dashboard_css() -> str:
    """Create custom CSS for dashboard styling"""
    return """
    /* Global styles */
    .gradio-container {
        background-color: #0f172a !important;  /* slate-900 */
    }
    
    /* Card styling */
    .dashboard-card {
        background-color: #1e293b !important;  /* slate-800 */
        border: 1px solid #334155 !important;  /* slate-700 */
        border-radius: 0.5rem !important;
        padding: 1rem !important;
        margin: 0.5rem 0 !important;
        color: #f1f5f9 !important;  /* slate-100 */
    }
    
    /* Chat container */
    .chatbot {
        background-color: #1e293b !important;  /* slate-800 */
        border-color: #334155 !important;  /* slate-700 */
    }
    
    /* Message bubbles */
    .chatbot .message.user {
        background-color: #334155 !important;  /* slate-700 */
        border: 1px solid #475569 !important;  /* slate-600 */
        color: #f1f5f9 !important;  /* slate-100 */
    }
    
    .chatbot .message.bot {
        background-color: #1e40af !important;  /* blue-800 */
        border: 1px solid #1e3a8a !important;  /* blue-900 */
        color: #f1f5f9 !important;  /* slate-100 */
    }
    
    /* Input fields */
    textarea, input[type="text"] {
        background-color: #334155 !important;  /* slate-700 */
        color: #f1f5f9 !important;  /* slate-100 */
        border: 1px solid #475569 !important;  /* slate-600 */
    }
    
    textarea:focus, input[type="text"]:focus {
        border-color: #3b82f6 !important;  /* blue-500 */
        box-shadow: 0 0 0 2px rgba(59, 130, 246, 0.2) !important;
    }
    
    /* Buttons */
    button.primary {
        background-color: #2563eb !important;  /* blue-600 */
        color: white !important;
    }
    
    button.primary:hover {
        background-color: #3b82f6 !important;  /* blue-500 */
    }
    
    button.secondary {
        background-color: #475569 !important;  /* slate-600 */
        color: white !important;
    }
    
    button.secondary:hover {
        background-color: #64748b !important;  /* slate-500 */
    }
    
    /* Tabs */
    .tab-nav {
        background-color: #1e293b !important;  /* slate-800 */
        border-bottom: 1px solid #334155 !important;  /* slate-700 */
    }
    
    .tab-nav button {
        color: #f1f5f9 !important;  /* slate-100 */
    }
    
    .tab-nav button.selected {
        border-bottom-color: #3b82f6 !important;  /* blue-500 */
    }
    
    /* Status indicators */
    .status-active {
        color: #22c55e !important;  /* green-500 */
        font-weight: 500 !important;
    }
    
    .status-completed {
        color: #94a3b8 !important;  /* slate-400 */
    }
    
    /* Headers */
    .dashboard-header {
        color: #f1f5f9 !important;  /* slate-100 */
        font-size: 1.5rem !important;
        font-weight: 600 !important;
        margin-bottom: 1rem !important;
    }
    
    /* Tables */
    table {
        background-color: #1e293b !important;  /* slate-800 */
        color: #f1f5f9 !important;  /* slate-100 */
    }
    
    th, td {
        border-color: #334155 !important;  /* slate-700 */
    }
    """

# %% ../nbs/02_learning_interface.ipynb 8
class LearningInterface:
    """
    Gradio interface for clinical learning interactions.
    
    Features:
    - Natural case discussion chat
    - Dynamic learning dashboard
    - Post-discussion analysis
    - Progress tracking
    """
    
    def __init__(
        self,
        context_path: Optional[Path] = None,
        theme: str = "default"
    ):
        """Initialize learning interface."""
        self.tutor = ClinicalTutor(context_path)
        self.theme = theme
        self.context_path = context_path
        
        # Track current discussion state
        self.current_discussion = {
            "started": None,
            "case_type": None,
            "messages": []
        }
        
        logger.info("Learning interface initialized")
    
    async def process_chat(
        self,
        message: str,
        history: List[List[str]],
        state: Dict[str, Any]
    ) -> Tuple[List[List[str]], str, Dict[str, Any]]: 
        """
        Process chat messages with state management.
        
        Args:
            message: User input message
            history: Chat history
            state: Current interface state
            
        Returns:
            tuple: (updated history, cleared message, updated state)
        """
        try:
            if not message.strip():
                return history, "", state
            
            # Start new discussion if none active
            if not state.get("discussion_active"):
                state["discussion_active"] = True
                state["discussion_start"] = datetime.now().isoformat()
            
            # Get tutor response
            response = await self.tutor.discuss_case(message)
            
            # Update history - now using list pairs instead of dicts
            if history is None:
                history = []
            history.append([message, response])  # Changed from dict format to list pair
            
            state["last_message"] = datetime.now().isoformat()
            
            return history, "", state
            
        except Exception as e:
            logger.error(f"Error in chat: {str(e)}")
            return history or [], "", state

    async def end_discussion(
        self,
        history: List[List[str]],
        state: Dict[str, Any]
    ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
        """
        Analyze completed discussion and prepare summary.
        
        Args:
            history: Chat history as list of [user_message, assistant_message] pairs
            state: Current interface state
            
        Returns:
            tuple: (analysis results, updated state)
        """
        try:
            if not history:
                return {
                    "learning_points": [],
                    "gaps": {},
                    "strengths": [],
                    "suggested_objectives": []
                }, state
            
            # Convert history format for analysis
            formatted_history = []
            for user_msg, assistant_msg in history:
                formatted_history.extend([
                    {"role": "user", "content": user_msg},
                    {"role": "assistant", "content": assistant_msg}
                ])
            
            # Get analysis
            analysis = await self.tutor.analyze_discussion(formatted_history)
            
            # Reset discussion state
            state["discussion_active"] = False
            state["discussion_start"] = None
            state["last_message"] = None
            
            return analysis, state
            
        except Exception as e:
            logger.error(f"Error analyzing discussion: {str(e)}")
            return {
                "learning_points": [],
                "gaps": {},
                "strengths": [],
                "suggested_objectives": []
            }, state    
            
    def update_rotation(
        self,
        specialty: str,
        start_date: str,
        end_date: str,
        focus_areas: str
    ) -> Tuple[str, str, str, str]:
        """
        Update rotation details and return updated values.
        
        Args:
            specialty: Rotation specialty
            start_date: Start date string
            end_date: End date string
            focus_areas: Comma-separated focus areas
            
        Returns:
            tuple: Updated field values
        """
        try:
            # Parse focus areas
            focus_list = [
                area.strip() 
                for area in focus_areas.split(",") 
                if area.strip()
            ]
            
            # Update context
            rotation = {
                "specialty": specialty,
                "start_date": start_date,
                "end_date": end_date,
                "key_focus_areas": focus_list
            }
            self.tutor.learning_context.update_rotation(rotation)
            
            # Return updated values
            return (
                specialty,
                start_date,
                end_date,
                ",".join(focus_list)
            )
            
        except Exception as e:
            logger.error(f"Error updating rotation: {str(e)}")
            current = self.tutor.learning_context.current_rotation
            return (
                current["specialty"],
                current["start_date"] or "",
                current["end_date"] or "",
                ",".join(current["key_focus_areas"])
            )

    def add_objective(
        self,
        objective: str,
        objectives_df: pd.DataFrame
    ) -> pd.DataFrame:
        """
        Add new learning objective and return updated dataframe.
        
        Args:
            objective: New objective text
            objectives_df: Current objectives dataframe
            
        Returns:
            pd.DataFrame: Updated objectives list
        """
        try:
            if not objective.strip():
                return objectives_df
                
            # Add to context
            self.tutor.learning_context.add_learning_objective(objective)
            
            # Convert to dataframe
            return pd.DataFrame([
                [obj["objective"], obj["status"], obj["added"]]
                for obj in self.tutor.learning_context.learning_objectives
            ], columns=["Objective", "Status", "Date Added"])
            
        except Exception as e:
            logger.error(f"Error adding objective: {str(e)}")
            return objectives_df

    def toggle_objective_status(
        self,
        evt: gr.SelectData,  # Updated to use gr.SelectData
        objectives_df: pd.DataFrame
    ) -> pd.DataFrame:
        """
        Toggle objective status between active and completed.
        
        Args:
            evt: Gradio select event containing row index
            objectives_df: Current objectives dataframe
            
        Returns:
            pd.DataFrame: Updated objectives list
        """
        try:
            objective_idx = evt.index[0]  # Get selected row index
            if objective_idx >= len(objectives_df):
                return objectives_df
                
            # Get objective
            objective = objectives_df.iloc[objective_idx]["Objective"]
            current_status = objectives_df.iloc[objective_idx]["Status"]
            
            # Toggle in context
            if current_status == "active":
                self.tutor.learning_context.complete_objective(objective)
            else:
                self.tutor.learning_context.add_learning_objective(objective)
            
            # Update dataframe
            return pd.DataFrame([
                [obj["objective"], obj["status"], obj["added"]]
                for obj in self.tutor.learning_context.learning_objectives
            ], columns=["Objective", "Status", "Date Added"])
            
        except Exception as e:
            logger.error(f"Error toggling objective: {str(e)}")
            return objectives_df

    def add_feedback_focus(
        self,
        focus: str,
        feedback_df: pd.DataFrame
    ) -> pd.DataFrame:
        """Add new feedback focus area."""
        try:
            if not focus.strip():
                return feedback_df
                
            # Add to context
            self.tutor.learning_context.toggle_feedback_focus(focus, True)
            
            # Update dataframe
            return pd.DataFrame([
                [pref["focus"], pref["active"]]
                for pref in self.tutor.learning_context.feedback_preferences
            ], columns=["Focus Area", "Active"])
            
        except Exception as e:
            logger.error(f"Error adding feedback focus: {str(e)}")
            return feedback_df

    def toggle_feedback_status(
        self,
        evt: gr.SelectData,  # Updated to use gr.SelectData
        feedback_df: pd.DataFrame
    ) -> pd.DataFrame:
        """Toggle feedback focus active status."""
        try:
            focus_idx = evt.index[0]  # Get selected row index
            if focus_idx >= len(feedback_df):
                return feedback_df
                
            # Get focus area
            focus = feedback_df.iloc[focus_idx]["Focus Area"]
            current_status = feedback_df.iloc[focus_idx]["Active"]
            
            # Toggle in context
            self.tutor.learning_context.toggle_feedback_focus(
                focus, 
                not current_status
            )
            
            # Update dataframe
            return pd.DataFrame([
                [pref["focus"], pref["active"]]
                for pref in self.tutor.learning_context.feedback_preferences
            ], columns=["Focus Area", "Active"])
            
        except Exception as e:
            logger.error(f"Error toggling feedback: {str(e)}")
            return feedback_df

    def create_interface(self) -> gr.Blocks:
        """Create and configure the Gradio interface"""
        with gr.Blocks(
            title="Clinical Learning Assistant",
            theme=self.theme,
            css=create_dashboard_css()
        ) as interface:
            # State management
            state = gr.State({
                "discussion_active": False,
                "discussion_start": None,
                "last_message": None
            })
            
            # Header
            with gr.Row():
                gr.Markdown(
                    "# Clinical Learning Assistant",
                    elem_classes=["dashboard-header"]
                )
            
            with gr.Row():
                # Left column - Chat interface
                with gr.Column(scale=2):
                    # Active discussion indicator
                    discussion_status = gr.Markdown(
                        "Start a new case discussion",
                        elem_classes=["dashboard-card"]
                    )
                    
                    # Chat interface
                    chatbot = gr.Chatbot(
                        height=500,
                        label="Case Discussion",
                        show_label=True,
                        elem_classes=["dashboard-card"]
                    )
                    
                    with gr.Row():
                        msg = gr.Textbox(
                            label="Present your case or ask questions",
                            placeholder=(
                                "Present your case as you would to your supervisor:\n"
                                "- Start with the chief complaint\n"
                                "- Include relevant history and findings\n"
                                "- Share your assessment and plan"
                            ),
                            lines=5
                        )
    
                        # Add voice input with updated syntax
                        audio_msg = gr.Audio(
                            label="Or speak your case",
                            sources=["microphone"],
                            type="numpy",
                            streaming=True
                        )
                    
                    with gr.Row():
                        clear = gr.Button("Clear Discussion")
                        end_discussion = gr.Button(
                            "End Discussion & Review",
                            variant="primary"
                        )
                
                # Right column - Learning dashboard
                with gr.Column(scale=1):
                    with gr.Tabs():
                        # Current Rotation tab
                        with gr.Tab("Current Rotation"):
                            with gr.Column(elem_classes=["dashboard-card"]):
                                specialty = gr.Textbox(
                                    label="Specialty",
                                    value=self.tutor.learning_context.current_rotation["specialty"]
                                )
                                start_date = gr.Textbox(
                                    label="Start Date (YYYY-MM-DD)",
                                    value=self.tutor.learning_context.current_rotation["start_date"]
                                )
                                end_date = gr.Textbox(
                                    label="End Date (YYYY-MM-DD)",
                                    value=self.tutor.learning_context.current_rotation["end_date"]
                                )
                                focus_areas = gr.Textbox(
                                    label="Key Focus Areas (comma-separated)",
                                    value=",".join(
                                        self.tutor.learning_context.current_rotation["key_focus_areas"]
                                    )
                                )
                                update_rotation_btn = gr.Button(
                                    "Update Rotation",
                                    variant="secondary"
                                )
                        
                        # Learning Objectives tab
                        with gr.Tab("Learning Objectives"):
                            with gr.Column(elem_classes=["dashboard-card"]):
                                objectives_df = gr.DataFrame(
                                    headers=["Objective", "Status", "Date Added"],
                                    value=[[
                                        obj["objective"],
                                        obj["status"],
                                        obj["added"]
                                    ] for obj in self.tutor.learning_context.learning_objectives],
                                    interactive=True,
                                    wrap=True
                                )
                                
                                with gr.Row():
                                    new_objective = gr.Textbox(
                                        label="New Learning Objective",
                                        placeholder="Enter objective..."
                                    )
                                    add_objective_btn = gr.Button(
                                        "Add",
                                        variant="secondary"
                                    )
                        
                        # Feedback Preferences tab
                        with gr.Tab("Feedback Focus"):
                            with gr.Column(elem_classes=["dashboard-card"]):
                                feedback_df = gr.DataFrame(
                                    headers=["Focus Area", "Active"],
                                    value=[[
                                        pref["focus"],
                                        pref["active"]
                                    ] for pref in self.tutor.learning_context.feedback_preferences],
                                    interactive=True,
                                    wrap=True
                                )
                                
                                with gr.Row():
                                    new_feedback = gr.Textbox(
                                        label="New Feedback Focus",
                                        placeholder="Enter focus area..."
                                    )
                                    add_feedback_btn = gr.Button(
                                        "Add",
                                        variant="secondary"
                                    )
                                    
                        # Knowledge Profile tab
                        with gr.Tab("Knowledge Profile"):
                            with gr.Column(elem_classes=["dashboard-card"]):
                                # Knowledge Gaps
                                gr.Markdown("### Knowledge Gaps")
                                gaps_display = gr.DataFrame(
                                    headers=["Topic", "Confidence"],
                                    value=[[
                                        topic, confidence
                                    ] for topic, confidence in 
                                        self.tutor.learning_context.knowledge_profile["gaps"].items()
                                    ],
                                    interactive=False
                                )
                                
                                # Strengths Display
                                gr.Markdown("### Strengths")
                                strengths_display = gr.DataFrame(
                                    headers=["Area"],
                                    value=[[strength] for strength in 
                                        self.tutor.learning_context.knowledge_profile["strengths"]
                                    ],
                                    interactive=False
                                )
                                
                                # Recent Progress
                                gr.Markdown("### Recent Progress")
                                progress_display = gr.DataFrame(
                                    headers=["Topic", "Improvement", "Date"],
                                    value=[[
                                        prog["topic"],
                                        f"{prog['improvement']:.2f}",
                                        prog["date"]
                                    ] for prog in 
                                        self.tutor.learning_context.knowledge_profile["recent_progress"]
                                    ],
                                    interactive=False
                                )
            
            # Discussion summary section
            summary_section = gr.Column(visible=False)
            with summary_section:
                gr.Markdown("## Discussion Summary")
                
                # Overview section
                with gr.Row():
                    with gr.Column():
                        gr.Markdown("### Session Overview")
                        session_overview = gr.JSON(
                            label="Discussion Details",
                            value={
                                "duration": "0 minutes",
                                "messages": 0,
                                "topics_covered": []
                            }
                        )
                
                # Learning Points and Gaps
                with gr.Row():
                    with gr.Column():
                        gr.Markdown("### Key Learning Points")
                        learning_points = gr.JSON(label="Points to Remember")
                    
                    with gr.Column():
                        gr.Markdown("### Knowledge Profile Updates")
                        with gr.Row():
                            gaps = gr.JSON(label="Areas for Improvement")
                            strengths = gr.JSON(label="Demonstrated Strengths")
                
                # Future Learning section
                gr.Markdown("### Planning Ahead")
                with gr.Row():
                    with gr.Column():
                        gr.Markdown("#### Suggested Learning Objectives")
                        objectives = gr.JSON(label="Consider Adding")
                    
                    with gr.Column():
                        gr.Markdown("#### Recommended Focus Areas")
                        recommendations = gr.JSON(label="Next Steps")
                
                # Action buttons
                with gr.Row():
                    add_selected_objectives = gr.Button(
                        "Add Selected Objectives",
                        variant="primary"
                    )
                    close_summary = gr.Button("Close Summary")
    
            # Event handlers
            # Add new event handler for voice input
            def process_audio(audio):
                if audio is None:
                    return None
                # Convert audio to text using your preferred method
                # For example, you could use transformers pipeline here
                try:
                    from transformers import pipeline
                    transcriber = pipeline("automatic-speech-recognition", model="openai/whisper-small")
                    text = transcriber(audio)["text"]
                    return text
                except Exception as e:
                    logger.error(f"Error transcribing audio: {str(e)}")
                    return None
            
            # Update the event handler:
            audio_msg.stop_recording(
                fn=process_audio,
                outputs=[msg]
            ).then(
                fn=self.process_chat,
                inputs=[msg, chatbot, state],
                outputs=[chatbot, msg, state]
            ).then(
                fn=self._update_discussion_status,
                inputs=[state],
                outputs=[discussion_status]
            )        

            msg.submit(
                self.process_chat,
                inputs=[msg, chatbot, state],
                outputs=[chatbot, msg, state]
            ).then(
                self._update_discussion_status,
                inputs=[state],
                outputs=[discussion_status]
            )
            
            clear.click(
                lambda: ([], "", {
                    "discussion_active": False,
                    "discussion_start": None,
                    "last_message": None
                }),
                outputs=[chatbot, msg, state]
            ).then(
                lambda: "Start a new case discussion",
                outputs=[discussion_status]
            )
            
            end_discussion.click(
                self.end_discussion,
                inputs=[chatbot, state],
                outputs=[
                    session_overview,
                    learning_points,
                    gaps,
                    strengths,
                    objectives,
                    recommendations
                ]
            ).then(
                lambda: gr.update(visible=True),
                None,
                summary_section
            ).then(
                self._refresh_knowledge_profile,
                outputs=[gaps_display, strengths_display, progress_display]
            )
            
            close_summary.click(
                lambda: gr.update(visible=False),
                None,
                summary_section
            )
            
            # Rotation management
            update_rotation_btn.click(
                self.update_rotation,
                inputs=[specialty, start_date, end_date, focus_areas],
                outputs=[specialty, start_date, end_date, focus_areas]
            )
            
            # Learning objectives management
            add_objective_btn.click(
                self.add_objective,
                inputs=[new_objective, objectives_df],
                outputs=[objectives_df]
            ).then(
                lambda: "",
                None,
                new_objective
            )
            
            objectives_df.select(
                self.toggle_objective_status,
                inputs=[objectives_df],
                outputs=[objectives_df]
            )
            
            # Feedback preferences management
            add_feedback_btn.click(
                self.add_feedback_focus,
                inputs=[new_feedback, feedback_df],
                outputs=[feedback_df]
            ).then(
                lambda: "",
                None,
                new_feedback
            )
            
            feedback_df.select(
                self.toggle_feedback_status,
                inputs=[feedback_df],
                outputs=[feedback_df]
            )
            
            # Add selected objectives from summary
            add_selected_objectives.click(
                self._add_suggested_objectives,
                inputs=[objectives],
                outputs=[objectives_df]
            )
    
            return interface
    
    def _update_discussion_status(self, state: Dict[str, Any]) -> str:
        """Update discussion status display"""
        try:
            if not state.get("discussion_active"):
                return "Start a new case discussion"
                
            start = datetime.fromisoformat(state["discussion_start"])
            duration = datetime.now() - start
            minutes = int(duration.total_seconds() / 60)
            
            return f"Active discussion ({minutes} minutes)"
            
        except Exception as e:
            logger.error(f"Error updating status: {str(e)}")
            return "Discussion status unknown"
    
    def _refresh_knowledge_profile(
        self
    ) -> Tuple[List[List[str]], List[List[str]], List[List[str]]]:
        """Refresh knowledge profile displays"""
        try:
            # Gaps
            gaps_data = [[
                topic, f"{confidence:.2f}"
            ] for topic, confidence in 
                self.tutor.learning_context.knowledge_profile["gaps"].items()
            ]
            
            # Strengths
            strengths_data = [[
                strength
            ] for strength in 
                self.tutor.learning_context.knowledge_profile["strengths"]
            ]
            
            # Progress
            progress_data = [[
                prog["topic"],
                f"{prog['improvement']:.2f}",
                prog["date"]
            ] for prog in 
                self.tutor.learning_context.knowledge_profile["recent_progress"]
            ]
            
            return gaps_data, strengths_data, progress_data
            
        except Exception as e:
            logger.error(f"Error refreshing profile: {str(e)}")
            return [], [], []
    
    def _add_suggested_objectives(
        self,
        evt: gr.SelectData,  # Updated to use gr.SelectData
        suggested_objectives: List[str]
    ) -> pd.DataFrame:
        """Add selected suggested objectives to learning objectives"""
        try:
            selected_indices = [evt.index[0]]  # Get selected row index
            
            for idx in selected_indices:
                if idx < len(suggested_objectives):
                    objective = suggested_objectives[idx]
                    self.tutor.learning_context.add_learning_objective(objective)
            
            return pd.DataFrame([
                [obj["objective"], obj["status"], obj["added"]]
                for obj in self.tutor.learning_context.learning_objectives
            ], columns=["Objective", "Status", "Date Added"])
            
        except Exception as e:
            logger.error(f"Error adding objectives: {str(e)}")
            return pd.DataFrame()

# %% ../nbs/02_learning_interface.ipynb 10
async def launch_learning_interface(
    port: Optional[int] = None,
    context_path: Optional[Path] = None,
    share: bool = False,
    theme: str = "default"
) -> None:
    """Launch the learning interface application."""
    try:
        interface = LearningInterface(context_path, theme)
        app = interface.create_interface()
        app.launch(
            server_port=port,
            share=share
        )
        logger.info(f"Interface launched on port: {port}")
    except Exception as e:
        logger.error(f"Error launching interface: {str(e)}")
        raise