streamlit_app.py

#!/usr/bin/env python3
"""
Clinical NLP Pipeline Streamlit App
Place this file at: clinical-nlp-pipeline/streamlit_app.py
"""

import streamlit as st
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import json
import sys
from pathlib import Path
from typing import List, Dict, Optional, Tuple
import re
from datetime import datetime
import time

# Add src to Python path
src_path = Path(__file__).parent / "src"
sys.path.insert(0, str(src_path))

# Import your existing modules
try:
    from clinical_nlp.nlp.entity_extraction import extract_clinical_entities
    from clinical_nlp.nlp.entity_extraction import Entity, EntityType
    from clinical_nlp.nlp.cpt_codes import CPTVectorStore, CPTReranker, CPTIntegration
except ImportError as e:
    st.error(f"Failed to import modules: {e}")
    st.error("Please ensure your modules are properly installed and accessible.")
    st.stop()

# Configure Streamlit page
st.set_page_config(
    page_title="Clinical NLP Pipeline",
    page_icon="🏥",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS
st.markdown("""
<style>
    .main-header {
        font-size: 3rem;
        color: #1f77b4;
        text-align: center;
        margin-bottom: 2rem;
    }
    
    .section-header {
        font-size: 1.5rem;
        color: #2e86ab;
        margin: 1.5rem 0 1rem 0;
        border-bottom: 2px solid #e0e0e0;
        padding-bottom: 0.5rem;
    }
    
    .entity-tag {
        display: inline-block;
        padding: 0.25rem 0.5rem;
        margin: 0.25rem;
        border-radius: 0.25rem;
        font-size: 0.875rem;
        font-weight: bold;
    }
    
    .entity-age { background-color: #ff9999; color: #000; }
    .entity-gender { background-color: #ffb3e6; color: #000; }
    .entity-condition { background-color: #ffcc99; color: #000; }
    .entity-symptom { background-color: #ffdb99; color: #000; }
    .entity-medication { background-color: #99ff99; color: #000; }
    .entity-dosage { background-color: #b3ffb3; color: #000; }
    .entity-vital-sign { background-color: #99ccff; color: #000; }
    .entity-lab-value { background-color: #ffff99; color: #000; }
    .entity-procedure { background-color: #cc99ff; color: #000; }
    .entity-test { background-color: #d9b3ff; color: #000; }
    .entity-body-part { background-color: #ffc266; color: #000; }
    .entity-date { background-color: #b3d9ff; color: #000; }
    .entity-duration { background-color: #c2f0fc; color: #000; }
    .entity-other { background-color: #e0e0e0; color: #000; }
    
    .cpt-card {
        border: 1px solid #ddd;
        border-radius: 8px;
        padding: 1rem;
        margin: 0.5rem 0;
        background-color: #f9f9f9;
    }
    
    .confidence-high { color: #28a745; font-weight: bold; }
    .confidence-medium { color: #ffc107; font-weight: bold; }
    .confidence-low { color: #dc3545; font-weight: bold; }
    
    .highlighted-text {
        padding: 1rem;
        border-radius: 8px;
        background-color: #f8f9fa;
        border: 1px solid #dee2e6;
        line-height: 1.8;
    }
    
    .metric-container {
        background-color: #f8f9fa;
        padding: 1rem;
        border-radius: 8px;
        border: 1px solid #dee2e6;
        margin: 0.5rem 0;
    }
</style>
""", unsafe_allow_html=True)

@st.cache_resource
def initialize_cpt_system():
    """Initialize CPT system components (cached for performance)"""
    try:
        # Initialize with your existing data directory
        data_dir = Path(__file__).parent / "cpt_data"
        vector_store = CPTVectorStore(data_dir=str(data_dir))
        reranker = CPTReranker()
        integration = CPTIntegration(vector_store, reranker)
        
        # Check if vector store exists, if not build it
        if not vector_store.cpt_codes:
            st.warning("CPT vector store not found. Please run the CPT system setup first.")
            return None
        
        return integration
    except Exception as e:
        st.error(f"Failed to initialize CPT system: {e}")
        return None

def get_entity_color_class(entity_type: EntityType) -> str:
    """Get CSS class for entity type"""
    color_map = {
        EntityType.AGE: "entity-age",
        EntityType.GENDER: "entity-gender", 
        EntityType.CONDITION: "entity-condition",
        EntityType.SYMPTOM: "entity-symptom",
        EntityType.MEDICATION: "entity-medication",
        EntityType.VITAL_SIGN: "entity-vital-sign",
        EntityType.LAB_VALUE: "entity-lab-value",
        EntityType.PROCEDURE: "entity-procedure",
        EntityType.TEST: "entity-test",
        EntityType.BODY_PART: "entity-body-part",
        EntityType.DATE: "entity-date",
        EntityType.DURATION: "entity-duration"
    }
    return color_map.get(entity_type, "entity-other")

def display_highlighted_text(text: str, entities: List[Entity]):
    """Display text with highlighted entities"""
    if not entities:
        st.markdown(f'<div class="highlighted-text">{text}</div>', unsafe_allow_html=True)
        return
    
    # Sort entities by start position
    sorted_entities = sorted(entities, key=lambda e: e.start)
    
    # Build highlighted text
    highlighted_text = ""
    last_end = 0
    
    for entity in sorted_entities:
        # Add text before entity
        highlighted_text += text[last_end:entity.start]
        
        # Add highlighted entity
        color_class = get_entity_color_class(entity.label)
        highlighted_text += f'<span class="entity-tag {color_class}" title="{entity.label.value} (conf: {entity.confidence:.2f})">{entity.text}</span>'
        
        last_end = entity.end
    
    # Add remaining text
    highlighted_text += text[last_end:]
    
    st.markdown(f'<div class="highlighted-text">{highlighted_text}</div>', unsafe_allow_html=True)

def display_entity_summary(entities: List[Entity]):
    """Display entity extraction summary"""
    if not entities:
        st.info("No entities extracted from the text.")
        return
    
    # Group entities by type
    entity_groups = {}
    for entity in entities:
        entity_type = entity.label.value
        if entity_type not in entity_groups:
            entity_groups[entity_type] = []
        entity_groups[entity_type].append(entity)
    
    # Create tabs for different views
    tab1, tab2, tab3 = st.tabs(["📊 Summary", "📋 Detailed List", "📈 Analytics"])
    
    with tab1:
        # Entity type summary
        cols = st.columns(4)
        
        with cols[0]:
            st.metric("Total Entities", len(entities))
        
        with cols[1]:
            st.metric("Entity Types", len(entity_groups))
        
        with cols[2]:
            avg_confidence = sum(e.confidence for e in entities) / len(entities)
            st.metric("Avg Confidence", f"{avg_confidence:.2f}")
        
        with cols[3]:
            high_conf_entities = sum(1 for e in entities if e.confidence > 0.8)
            st.metric("High Confidence", f"{high_conf_entities}/{len(entities)}")
        
        # Entity type breakdown
        st.subheader("Entity Type Breakdown")
        entity_counts = {k: len(v) for k, v in entity_groups.items()}
        
        if entity_counts:
            fig = px.bar(
                x=list(entity_counts.keys()),
                y=list(entity_counts.values()),
                title="Entities by Type",
                labels={'x': 'Entity Type', 'y': 'Count'}
            )
            fig.update_layout(height=400)
            st.plotly_chart(fig, use_container_width=True)
    
    with tab2:
        # Detailed entity list
        for entity_type, entities_of_type in entity_groups.items():
            with st.expander(f"{entity_type} ({len(entities_of_type)} items)"):
                for i, entity in enumerate(entities_of_type, 1):
                    col1, col2, col3, col4 = st.columns([3, 1, 1, 2])
                    
                    with col1:
                        st.write(f"**{entity.text}**")
                    
                    with col2:
                        confidence_class = "confidence-high" if entity.confidence > 0.8 else "confidence-medium" if entity.confidence > 0.6 else "confidence-low"
                        st.markdown(f'<span class="{confidence_class}">{entity.confidence:.2f}</span>', unsafe_allow_html=True)
                    
                    with col3:
                        st.write(f"[{entity.start}-{entity.end}]")
                    
                    with col4:
                        source = getattr(entity, 'source', 'unknown')
                        st.write(source)
                    
                    # Show metadata if available
                    if entity.metadata:
                        st.json(entity.metadata, expanded=False)
    
    with tab3:
        # Analytics
        col1, col2 = st.columns(2)
        
        with col1:
            # Confidence distribution
            confidences = [e.confidence for e in entities]
            fig = px.histogram(
                x=confidences,
                nbins=20,
                title="Confidence Score Distribution",
                labels={'x': 'Confidence Score', 'y': 'Count'}
            )
            st.plotly_chart(fig, use_container_width=True)
        
        with col2:
            # Source breakdown
            sources = [getattr(e, 'source', 'unknown') for e in entities]
            source_counts = pd.Series(sources).value_counts()
            
            fig = px.pie(
                values=source_counts.values,
                names=source_counts.index,
                title="Entities by Source"
            )
            st.plotly_chart(fig, use_container_width=True)

def display_cpt_suggestions(suggestions: Dict):
    """Display CPT code suggestions"""
    if not suggestions or not suggestions.get('suggestions'):
        st.info("No CPT code suggestions generated.")
        return
    
    # Summary metrics
    col1, col2, col3, col4 = st.columns(4)
    
    with col1:
        st.metric("Suggested Codes", len(suggestions['suggestions']))
    
    with col2:
        total_rvu = suggestions.get('total_rvu', 0)
        st.metric("Total RVU", f"{total_rvu:.2f}")
    
    with col3:
        avg_confidence = sum(s['confidence'] for s in suggestions['suggestions']) / len(suggestions['suggestions'])
        st.metric("Avg Confidence", f"{avg_confidence:.2f}")
    
    with col4:
        est_revenue = total_rvu * 35.89  # 2024 conversion factor
        st.metric("Est. Revenue", f"${est_revenue:.2f}")
    
    # Display search context
    with st.expander("🔍 Search Context"):
        col1, col2 = st.columns(2)
        
        with col1:
            st.write("**Generated Query:**")
            st.code(suggestions.get('query', 'N/A'))
        
        with col2:
            st.write("**Clinical Context:**")
            context = suggestions.get('context', {})
            for key, value in context.items():
                if value != 'unknown':
                    st.write(f"• **{key.replace('_', ' ').title()}:** {value}")
    
    # CPT Code Cards
    st.subheader("🏥 Suggested CPT Codes")
    
    for i, suggestion in enumerate(suggestions['suggestions'], 1):
        with st.container():
            st.markdown(f'<div class="cpt-card">', unsafe_allow_html=True)
            
            # Header row
            col1, col2, col3, col4 = st.columns([2, 1, 1, 1])
            
            with col1:
                st.markdown(f"### {i}. CPT {suggestion['code']}")
                st.write(suggestion['description'])
            
            with col2:
                confidence = suggestion['confidence']
                confidence_class = "confidence-high" if confidence > 0.8 else "confidence-medium" if confidence > 0.6 else "confidence-low"
                st.markdown(f'**Confidence:** <span class="{confidence_class}">{confidence:.2f}</span>', unsafe_allow_html=True)
            
            with col3:
                st.markdown(f"**RVU:** {suggestion['rvu']:.2f}")
                st.markdown(f"**Category:** {suggestion['category']}")
            
            with col4:
                revenue = suggestion['rvu'] * 35.89
                st.markdown(f"**Revenue:** ${revenue:.2f}")
            
            # Reasons
            if suggestion.get('reasons'):
                st.markdown("**Justification:**")
                for reason in suggestion['reasons']:
                    st.write(f"• {reason}")
            
            st.markdown('</div>', unsafe_allow_html=True)
            st.markdown("---")

def get_sample_texts() -> Dict[str, str]:
    """Get sample clinical texts"""
    return {
        "Custom Text": "",
        
        "Cardiology Visit": """Patient is a 65-year-old male with history of hypertension and diabetes mellitus who presented with chest pain and shortness of breath. Vital signs: BP 140/90, HR 78, RR 18, O2 sat 96%, temp 98.6°F. Troponin peaked at 45.2 ng/mL. Echocardiogram showed reduced ejection fraction. Started on aspirin 81mg daily, lisinopril 10mg daily, and metoprolol 25mg twice daily. Patient underwent cardiac catheterization which revealed 90% LAD stenosis.""",
        
        "Routine Office Visit": """45-year-old established patient presenting for routine follow-up of hypertension and diabetes. Blood pressure today 135/85. HbA1c 7.2%. ECG performed showing normal sinus rhythm. Patient reports good medication compliance. Continue metformin 500mg twice daily and lisinopril 10mg daily. Follow up in 6 months.""",
        
        "Hospital Admission": """78-year-old female admitted with community-acquired pneumonia. Patient presented with 3-day history of fever, productive cough, and dyspnea. Vital signs: T 101.2°F, BP 130/85, HR 95, RR 22, O2 sat 88%. Chest X-ray shows right lower lobe infiltrate. WBC count 15,000. Started on ceftriaxone 1g IV daily and azithromycin 500mg daily. Oxygen therapy at 2L/min.""",
        
        "Emergency Visit": """32-year-old male brought to ED after motor vehicle accident. Complaint of severe right leg pain, unable to bear weight. Alert and oriented x3. Vital signs stable: BP 135/82, HR 95, RR 18, temp 98.7°F. Right tibia shows obvious deformity with swelling. X-ray confirms displaced tibial fracture. Morphine 4mg IV for pain control. Orthopedic surgery consulted for ORIF.""",
        
        "Diabetes Follow-up": """28-year-old type 1 diabetic presents for routine follow-up. HbA1c today 8.2%, up from 7.1% three months ago. Patient admits to missed insulin doses due to work schedule. No symptoms of hyperglycemia currently. Blood pressure 118/75, weight stable. Discussed importance of medication compliance. Adjusted insulin regimen and scheduled diabetes educator visit."""
    }

def main():
    """Main Streamlit application"""
    
    # Header
    st.markdown('<h1 class="main-header">🏥 Clinical NLP Pipeline</h1>', unsafe_allow_html=True)
    st.markdown("**Extract clinical entities and suggest CPT codes from medical text**")
    
    # Sidebar configuration
    with st.sidebar:
        st.header("⚙️ Configuration")
        
        # Processing options
        st.subheader("Processing Options")
        extract_entities_enabled = st.checkbox("Extract Entities", value=True)
        suggest_cpt_enabled = st.checkbox("Suggest CPT Codes", value=True)
        
        # CPT options
        if suggest_cpt_enabled:
            st.subheader("CPT Options")
            top_k_cpt = st.slider("Number of CPT suggestions", 1, 10, 5)
        
        # Sample texts
        st.subheader("📝 Sample Texts")
        sample_options = get_sample_texts()
        selected_sample = st.selectbox("Choose sample text:", list(sample_options.keys()))
        
        # System info
        st.subheader("ℹ️ System Info")
        st.info(f"**Data Directory:** {Path(__file__).parent / 'cpt_data'}")
        
        # Check if CPT system is available
        if suggest_cpt_enabled:
            cpt_available = Path(__file__).parent / "cpt_data" / "vector_store" / "cpt_codes.json"
            if cpt_available.exists():
                st.success("✅ CPT Vector Store Ready")
            else:
                st.warning("⚠️ CPT Vector Store Not Built")
                st.markdown("Run the CPT system setup first")
    
    # Main content area
    col1, col2 = st.columns([2, 1])
    
    with col1:
        st.subheader("📝 Clinical Text Input")
        
        # Text input
        if selected_sample != "Custom Text":
            default_text = sample_options[selected_sample]
        else:
            default_text = st.session_state.get('clinical_text', '')
        
        clinical_text = st.text_area(
            "Enter clinical text:",
            value=default_text,
            height=300,
            placeholder="Enter clinical notes, discharge summaries, or other medical text..."
        )
        
        # Process button
        process_button = st.button("🚀 Process Text", type="primary", use_container_width=True)
    
    with col2:
        st.subheader("📊 Text Statistics")
        if clinical_text:
            st.markdown('<div class="metric-container">', unsafe_allow_html=True)
            st.metric("Characters", f"{len(clinical_text):,}")
            st.metric("Words", f"{len(clinical_text.split()):,}")
            sentences = len([s for s in clinical_text.split('.') if s.strip()])
            st.metric("Sentences", f"{sentences}")
            
            # Text complexity indicators
            avg_word_length = sum(len(word) for word in clinical_text.split()) / max(len(clinical_text.split()), 1)
            st.metric("Avg Word Length", f"{avg_word_length:.1f}")
            st.markdown('</div>', unsafe_allow_html=True)
        else:
            st.info("Enter text to see statistics")
    
    # Process the text when button is clicked
    if process_button and clinical_text.strip():
        
        # Store text in session state
        st.session_state['clinical_text'] = clinical_text
        
        # Initialize components
        cpt_integration = None
        if suggest_cpt_enabled:
            with st.spinner("Initializing CPT system..."):
                cpt_integration = initialize_cpt_system()
        
        if extract_entities_enabled or (suggest_cpt_enabled and cpt_integration):
            
            # Processing section
            st.markdown('<h2 class="section-header">🔄 Processing Results</h2>', unsafe_allow_html=True)
            
            # Create progress bar
            progress_bar = st.progress(0)
            status_text = st.empty()
            
            # Start processing
            start_time = time.time()
            entities = []
            suggestions = {}
            
            # Step 1: Entity extraction
            if extract_entities_enabled:
                status_text.text("🔍 Extracting clinical entities...")
                progress_bar.progress(25)
                
                try:
                    ner_start = time.time()
                    entities = extract_clinical_entities(clinical_text)
                    ner_time = time.time() - ner_start
                    st.session_state['ner_time'] = ner_time
                    status_text.text(f"✅ Extracted {len(entities)} entities in {ner_time:.2f}s")
                except Exception as e:
                    st.error(f"Entity extraction failed: {e}")
                    entities = []
                
                progress_bar.progress(50)
            
            # Step 2: CPT code suggestions
            if suggest_cpt_enabled and cpt_integration:
                status_text.text("🏥 Generating CPT code suggestions...")
                progress_bar.progress(75)
                
                try:
                    search_start = time.time()
                    suggestions = cpt_integration.suggest_cpt_codes(
                        clinical_text, 
                        entities, 
                        top_k=top_k_cpt
                    )
                    search_time = time.time() - search_start
                    st.session_state['search_time'] = search_time
                    status_text.text(f"✅ Generated {len(suggestions.get('suggestions', []))} CPT suggestions in {search_time:.2f}s")
                except Exception as e:
                    st.error(f"CPT suggestion failed: {e}")
                    suggestions = {}
            
            # Complete processing
            total_time = time.time() - start_time
            st.session_state['processing_time'] = total_time
            
            progress_bar.progress(100)
            status_text.text(f"🎉 Processing complete! Total time: {total_time:.2f}s")
            
            # Clear progress indicators after a moment
            time.sleep(1)
            progress_bar.empty()
            status_text.empty()
            
            # Display results
            if extract_entities_enabled and entities:
                st.markdown('<h2 class="section-header">🏷️ Extracted Entities</h2>', unsafe_allow_html=True)
                
                # Highlighted text
                st.subheader("📄 Text with Highlighted Entities")
                display_highlighted_text(clinical_text, entities)
                
                # Entity details
                st.subheader("📋 Entity Analysis")
                display_entity_summary(entities)
            
            if suggest_cpt_enabled and suggestions:
                st.markdown('<h2 class="section-header">🏥 CPT Code Suggestions</h2>', unsafe_allow_html=True)
                display_cpt_suggestions(suggestions)
            
            # Export section
            if entities or suggestions:
                st.markdown('<h2 class="section-header">💾 Export Results</h2>', unsafe_allow_html=True)
                
                col1, col2, col3 = st.columns(3)
                
                with col1:
                    if entities:
                        entities_json = json.dumps([
                            {
                                'text': e.text,
                                'label': e.label.value,
                                'start': e.start,
                                'end': e.end,
                                'confidence': e.confidence,
                                'source': getattr(e, 'source', 'unknown'),
                                'metadata': getattr(e, 'metadata', {})
                            } for e in entities
                        ], indent=2)
                        
                        st.download_button(
                            "📥 Download Entities",
                            entities_json,
                            f"entities_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json",
                            "application/json",
                            use_container_width=True
                        )
                
                with col2:
                    if suggestions:
                        suggestions_json = json.dumps(suggestions, indent=2)
                        
                        st.download_button(
                            "📥 Download CPT Codes",
                            suggestions_json,
                            f"cpt_suggestions_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json",
                            "application/json",
                            use_container_width=True
                        )
                
                with col3:
                    # Create comprehensive report
                    report_data = {
                        'timestamp': datetime.now().isoformat(),
                        'clinical_text': clinical_text,
                        'processing_time': total_time,
                        'entities': json.loads(entities_json) if entities else None,
                        'cpt_suggestions': suggestions
                    }
                    
                    report_json = json.dumps(report_data, indent=2, default=str)
                    
                    st.download_button(
                        "📥 Download Full Report",
                        report_json,
                        f"clinical_nlp_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json",
                        "application/json",
                        use_container_width=True
                    )
        
        else:
            st.warning("Please enable at least one processing option (Entity Extraction or CPT Suggestions)")
    
    elif process_button:
        st.warning("Please enter some clinical text to process.")
    
    # Footer
    st.markdown("---")
    col1, col2, col3 = st.columns([1, 2, 1])
    with col2:
        st.markdown(
            "<div style='text-align: center'>"
            "Built with ❤️ using Streamlit | Clinical NLP Pipeline v1.0"
            "</div>", 
            unsafe_allow_html=True
        )

if __name__ == "__main__":
    main()