app.py

import streamlit as st
import pandas as pd
import numpy as np
import joblib
import pickle
from PIL import Image
import io
import cv2
import pytesseract
from sklearn.metrics import roc_auc_score, accuracy_score, classification_report
import plotly.graph_objects as go
import plotly.express as px
from datetime import datetime
import requests
import json
import base64

# Custom CSS for styling
def local_css():
    st.markdown("""
    <style>
    .main-header {
        font-size: 3rem;
        color: #2E86AB;
        text-align: center;
        margin-bottom: 2rem;
        font-weight: bold;
    }
    .urdu-text {
        font-family: 'Jameel Noori Nastaleeq', 'Noto Sans Arabic';
        font-size: 1.2rem;
        direction: rtl;
    }
    .risk-high { background-color: #ffcccc; padding: 10px; border-radius: 5px; }
    .risk-medium { background-color: #fff3cd; padding: 10px; border-radius: 5px; }
    .risk-low { background-color: #d4edda; padding: 10px; border-radius: 5px; }
    .priority-box { 
        border: 2px solid #2E86AB; 
        padding: 20px; 
        border-radius: 10px; 
        margin: 10px 0;
    }
    </style>
    """, unsafe_allow_html=True)

# Initialize session state
def init_session_state():
    if 'language' not in st.session_state:
        st.session_state.language = 'English'
    if 'patient_data' not in st.session_state:
        st.session_state.patient_data = {}
    if 'risk_scores' not in st.session_state:
        st.session_state.risk_scores = {}

# Load models with error handling
@st.cache_resource
def load_models():
    try:
        heart_model = joblib.load("best_model.pkl")
        diabetes_model = joblib.load("best_diabetes_model.pkl")
        hypertension_model = joblib.load("hypertension_model.pkl")
        return heart_model, diabetes_model, hypertension_model
    except Exception as e:
        st.error(f"Error loading models: {str(e)}")
        return None, None, None

# Urdu translations
URDU_TRANSLATIONS = {
    "AI-Priority OPD System": "AI-ترجیحی OPD سسٹم",
    "Patient Information": "مریض کی معلومات",
    "Name": "نام",
    "Age": "عمر",
    "Gender": "جنس",
    "Contact": "رابطہ نمبر",
    "Medical History": "طبی تاریخ",
    "Vital Signs": "اہم علامات",
    "Blood Pressure (systolic)": "بلڈ پریشر (سسٹولک)",
    "Blood Pressure (diastolic)": "بلڈ پریشر ڈائیسٹولک)",
    "Heart Rate": "دل کی دھڑکن",
    "Cholesterol Level": "کولیسٹرول کی سطح",
    "Blood Glucose": "خون میں گلوکوز",
    "BMI": "باڈی ماس انڈیکس",
    "Symptoms": "علامات",
    "Chest Pain": "سینے میں درد",
    "Shortness of Breath": "سانس لینے میں دشواری",
    "Fatigue": "تھکاوٹ",
    "Upload Prescription": "نسخہ اپ لوڈ کریں",
    "Calculate Risk Score": "خطرے کا اسکور معلوم کریں",
    "High Priority - Emergency Care Required": "اعلی ترجیح - ہنگامی علاج کی ضرورت",
    "Medium Priority - Same Day Consultation": "درمیانی ترجیح - اسی دن مشورہ",
    "Low Priority - Routine Appointment": "کم ترجیح - معمول کی ملاقات",
    "Healthcare Chatbot": "ہیلتھ کیئر چیٹ بوٹ",
    "Ask health-related questions": "صحت سے متعلق سوالات پوچھیں"
}

class OCRProcessor:
    def __init__(self):
        pass
    
    def preprocess_image(self, image):
        """Enhance image for better OCR accuracy"""
        try:
            # Convert to grayscale
            gray = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
            
            # Apply thresholding
            _, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
            
            # Remove noise
            kernel = np.ones((1, 1), np.uint8)
            processed = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel)
            
            return processed
        except Exception as e:
            st.error(f"Image processing error: {str(e)}")
            return np.array(image)
    
    def extract_text(self, image):
        """Extract text from prescription image"""
        try:
            processed_image = self.preprocess_image(image)
            
            # Configure Tesseract for medical text
            custom_config = r'--oem 3 --psm 6 -l eng'
            text = pytesseract.image_to_string(processed_image, config=custom_config)
            
            return text.strip()
        except Exception as e:
            st.error(f"OCR Error: {str(e)}")
            return ""

class HealthcareChatbot:
    def __init__(self):
        self.health_tips = {
            'heart': [
                "Maintain a healthy diet low in saturated fats",
                "Exercise regularly for at least 30 minutes daily",
                "Monitor blood pressure regularly",
                "Avoid smoking and limit alcohol consumption"
            ],
            'diabetes': [
                "Monitor blood sugar levels regularly",
                "Follow a balanced diet with controlled carbohydrates",
                "Take medications as prescribed",
                "Stay physically active"
            ],
            'hypertension': [
                "Reduce salt intake in your diet",
                "Practice stress management techniques",
                "Maintain healthy body weight",
                "Limit caffeine consumption"
            ]
        }
    
    def get_response(self, query, language='English'):
        """Generate chatbot response"""
        query_lower = query.lower()
        
        # Simple rule-based responses
        if any(word in query_lower for word in ['heart', 'cardiac', 'chest pain']):
            tips = self.health_tips['heart']
        elif any(word in query_lower for word in ['diabetes', 'sugar', 'glucose']):
            tips = self.health_tips['diabetes']
        elif any(word in query_lower for word in ['blood pressure', 'hypertension']):
            tips = self.health_tips['hypertension']
        else:
            tips = ["Maintain regular health checkups", "Follow your doctor's advice", "Stay hydrated and eat balanced meals"]
        
        response = f"Based on your query, here are some health tips:\n" + "\n".join([f"• {tip}" for tip in tips])
        
        if language == 'Urdu':
            # Simple Urdu translations (in a real app, use proper translation API)
            response = "آپ کے سوال کے مطابق، یہاں کچھ صحت کے نکات ہیں:\n" + "\n".join([f"• {tip}" for tip in tips])
        
        return response

def calculate_priority_score(heart_risk, diabetes_risk, hypertension_risk):
    """Calculate integrated priority score"""
    # Weighted priority scoring based on clinical severity
    priority = (
        heart_risk * 0.4 + 
        diabetes_risk * 0.3 + 
        hypertension_risk * 0.3
    )
    
    return priority

def get_priority_recommendation(priority_score, language='English'):
    """Get priority-based recommendation"""
    if priority_score > 0.8:
        if language == 'Urdu':
            return "EMERGENCY_CARE", "اعلی ترجیح - ہنگامی علاج کی ضرورت", "risk-high"
        else:
            return "EMERGENCY_CARE", "High Priority - Emergency Care Required", "risk-high"
    elif priority_score > 0.6:
        if language == 'Urdu':
            return "SAME_DAY_CONSULT", "درمیانی ترجیح - اسی دن مشورہ", "risk-medium"
        else:
            return "SAME_DAY_CONSULT", "Medium Priority - Same Day Consultation", "risk-medium"
    else:
        if language == 'Urdu':
            return "ROUTINE_APPOINTMENT", "کم ترجیح - معمول کی ملاقات", "risk-low"
        else:
            return "ROUTINE_APPOINTMENT", "Low Priority - Routine Appointment", "risk-low"

def main():
    # Page configuration
    st.set_page_config(
        page_title="AI-Priority OPD System",
        page_icon="🏥",
        layout="wide",
        initial_sidebar_state="expanded"
    )
    
    # Load custom CSS
    local_css()
    init_session_state()
    
    # Load models
    heart_model, diabetes_model, hypertension_model = load_models()
    
    if heart_model is None or diabetes_model is None or hypertension_model is None:
        st.error("Failed to load ML models. Please check model files.")
        return
    
    # Initialize processors
    ocr_processor = OCRProcessor()
    chatbot = HealthcareChatbot()
    
    # Language selector in sidebar
    with st.sidebar:
        st.image("https://via.placeholder.com/150x50/2E86AB/FFFFFF?text=AI-OPD", use_column_width=True)
        language = st.radio("Select Language / زبان منتخب کریں", 
                           ["English", "Urdu"], 
                           key="language_selector")
        
        st.markdown("---")
        st.subheader("Quick Actions")
        if st.button("New Patient Assessment" if language == "English" else "نیا مریض تشخیص"):
            st.session_state.patient_data = {}
            st.session_state.risk_scores = {}
    
    # Main header
    if language == "English":
        st.markdown('<h1 class="main-header">🏥 AI-Priority OPD System</h1>', unsafe_allow_html=True)
        st.markdown("### Smart Patient Triage and Priority Management")
    else:
        st.markdown('<h1 class="main-header">🏥 AI-ترجیحی OPD سسٹم</h1>', unsafe_allow_html=True)
        st.markdown("### ذہین مریض کی درجہ بندی اور ترجیحی انتظام")
    
    # Create tabs
    tab1, tab2, tab3, tab4 = st.tabs([
        "Patient Assessment" if language == "English" else "مریض تشخیص",
        "Prescription OCR" if language == "English" else "نسخہ OCR",
        "Health Chatbot" if language == "English" else "ہیلتھ چیٹ بوٹ",
        "Analytics" if language == "English" else "تجزیات"
    ])
    
    with tab1:
        # Patient Assessment Form
        if language == "English":
            st.header("Patient Information")
        else:
            st.header("مریض کی معلومات")
        
        col1, col2 = st.columns(2)
        
        with col1:
            # Basic Information
            if language == "English":
                name = st.text_input("Full Name")
                age = st.number_input("Age", min_value=0, max_value=120, value=30)
                gender = st.selectbox("Gender", ["Male", "Female", "Other"])
                contact = st.text_input("Contact Number")
            else:
                name = st.text_input("مکمل نام")
                age = st.number_input("عمر", min_value=0, max_value=120, value=30)
                gender = st.selectbox("جنس", ["مرد", "عورت", "دیگر"])
                contact = st.text_input("رابطہ نمبر")
        
        with col2:
            # Vital Signs
            if language == "English":
                st.subheader("Vital Signs")
                bp_systolic = st.number_input("Blood Pressure (systolic)", min_value=50, max_value=250, value=120)
                bp_diastolic = st.number_input("Blood Pressure (diastolic)", min_value=30, max_value=150, value=80)
                heart_rate = st.number_input("Heart Rate (bpm)", min_value=30, max_value=200, value=72)
                cholesterol = st.number_input("Cholesterol Level (mg/dL)", min_value=100, max_value=400, value=200)
                glucose = st.number_input("Blood Glucose (mg/dL)", min_value=50, max_value=500, value=100)
                bmi = st.number_input("BMI", min_value=10.0, max_value=50.0, value=22.0, step=0.1)
            else:
                st.subheader("اہم علامات")
                bp_systolic = st.number_input("بلڈ پریشر (سسٹولک)", min_value=50, max_value=250, value=120)
                bp_diastolic = st.number_input("بلڈ پریشر (ڈائیسٹولک)", min_value=30, max_value=150, value=80)
                heart_rate = st.number_input("دل کی دھڑکن (bpm)", min_value=30, max_value=200, value=72)
                cholesterol = st.number_input("کولیسٹرول کی سطح (mg/dL)", min_value=100, max_value=400, value=200)
                glucose = st.number_input("خون میں گلوکوز (mg/dL)", min_value=50, max_value=500, value=100)
                bmi = st.number_input("باڈی ماس انڈیکس", min_value=10.0, max_value=50.0, value=22.0, step=0.1)
        
        # Symptoms
        if language == "English":
            st.subheader("Symptoms")
            col3, col4 = st.columns(2)
            with col3:
                chest_pain = st.checkbox("Chest Pain")
                shortness_breath = st.checkbox("Shortness of Breath")
            with col4:
                fatigue = st.checkbox("Fatigue")
                dizziness = st.checkbox("Dizziness")
        else:
            st.subheader("علامات")
            col3, col4 = st.columns(2)
            with col3:
                chest_pain = st.checkbox("سینے میں درد")
                shortness_breath = st.checkbox("سانس لینے میں دشواری")
            with col4:
                fatigue = st.checkbox("تھکاوٹ")
                dizziness = st.checkbox("چکر آنا")
        
        # Risk Assessment Button
        if language == "English":
            assess_button = st.button("🚀 Calculate Risk Score & Priority", use_container_width=True)
        else:
            assess_button = st.button("🚀 خطرے کا اسکور اور ترجیح معلوم کریں", use_container_width=True)
        
        if assess_button:
            try:
                # Prepare feature arrays (adjust based on your model requirements)
                heart_features = np.array([[age, bp_systolic, cholesterol, heart_rate, 1 if chest_pain else 0]])
                diabetes_features = np.array([[age, glucose, bmi, cholesterol]])
                hypertension_features = np.array([[age, bp_systolic, bp_diastolic, bmi]])
                
                # Get predictions
                heart_risk = heart_model.predict_proba(heart_features)[0][1]
                diabetes_risk = diabetes_model.predict_proba(diabetes_features)[0][1]
                hypertension_risk = hypertension_model.predict_proba(hypertension_features)[0][1]
                
                # Calculate priority score
                priority_score = calculate_priority_score(heart_risk, diabetes_risk, hypertension_risk)
                priority_level, recommendation, risk_class = get_priority_recommendation(priority_score, language)
                
                # Store results
                st.session_state.risk_scores = {
                    'heart': heart_risk,
                    'diabetes': diabetes_risk,
                    'hypertension': hypertension_risk,
                    'priority': priority_score,
                    'recommendation': recommendation,
                    'level': priority_level
                }
                
                # Display results
                st.markdown("---")
                
                # Risk Scores Visualization
                col5, col6, col7, col8 = st.columns(4)
                
                with col5:
                    fig = go.Figure(go.Indicator(
                        mode = "gauge+number+delta",
                        value = heart_risk,
                        domain = {'x': [0, 1], 'y': [0, 1]},
                        title = {'text': "Heart Disease Risk"},
                        gauge = {'axis': {'range': [0, 1]},
                                'bar': {'color': "red"},
                                'steps': [{'range': [0, 0.3], 'color': "lightgreen"},
                                         {'range': [0.3, 0.7], 'color': "yellow"},
                                         {'range': [0.7, 1], 'color': "red"}]}))
                    st.plotly_chart(fig, use_container_width=True)
                
                with col6:
                    fig = go.Figure(go.Indicator(
                        mode = "gauge+number+delta",
                        value = diabetes_risk,
                        domain = {'x': [0, 1], 'y': [0, 1]},
                        title = {'text': "Diabetes Risk"},
                        gauge = {'axis': {'range': [0, 1]},
                                'bar': {'color': "orange"},
                                'steps': [{'range': [0, 0.3], 'color': "lightgreen"},
                                         {'range': [0.3, 0.7], 'color': "yellow"},
                                         {'range': [0.7, 1], 'color': "red"}]}))
                    st.plotly_chart(fig, use_container_width=True)
                
                with col7:
                    fig = go.Figure(go.Indicator(
                        mode = "gauge+number+delta",
                        value = hypertension_risk,
                        domain = {'x': [0, 1], 'y': [0, 1]},
                        title = {'text': "Hypertension Risk"},
                        gauge = {'axis': {'range': [0, 1]},
                                'bar': {'color': "blue"},
                                'steps': [{'range': [0, 0.3], 'color': "lightgreen"},
                                         {'range': [0.3, 0.7], 'color': "yellow"},
                                         {'range': [0.7, 1], 'color': "red"}]}))
                    st.plotly_chart(fig, use_container_width=True)
                
                with col8:
                    fig = go.Figure(go.Indicator(
                        mode = "gauge+number+delta",
                        value = priority_score,
                        domain = {'x': [0, 1], 'y': [0, 1]},
                        title = {'text': "Priority Score"},
                        gauge = {'axis': {'range': [0, 1]},
                                'bar': {'color': "purple"},
                                'steps': [{'range': [0, 0.6], 'color': "lightgreen"},
                                         {'range': [0.6, 0.8], 'color': "yellow"},
                                         {'range': [0.8, 1], 'color': "red"}]}))
                    st.plotly_chart(fig, use_container_width=True)
                
                # Priority Recommendation
                st.markdown(f'<div class="priority-box {risk_class}">', unsafe_allow_html=True)
                if language == "English":
                    st.markdown(f"## 🎯 Priority Recommendation: {recommendation}")
                    st.markdown(f"**Overall Risk Score:** {priority_score:.3f}")
                    st.markdown(f"**Recommended Action:** {priority_level.replace('_', ' ').title()}")
                else:
                    st.markdown(f"## 🎯 ترجیحی سفارش: {recommendation}")
                    st.markdown(f"**کل خطرے کا اسکور:** {priority_score:.3f}")
                    st.markdown(f"**سفارش کردہ عمل:** {priority_level.replace('_', ' ').title()}")
                st.markdown('</div>', unsafe_allow_html=True)
                
            except Exception as e:
                st.error(f"Error in risk assessment: {str(e)}")
    
    with tab2:
        # Prescription OCR
        if language == "English":
            st.header("Prescription OCR Analysis")
            st.write("Upload a prescription image to extract medication information")
        else:
            st.header("نسخہ OCR تجزیہ")
            st.write("دوائی کی معلومات نکالنے کے لیے نسخہ کی تصویر اپ لوڈ کریں")
        
        uploaded_file = st.file_uploader(
            "Upload Prescription Image" if language == "English" else "نسخہ تصویر اپ لوڈ کریں",
            type=['png', 'jpg', 'jpeg']
        )
        
        if uploaded_file is not None:
            image = Image.open(uploaded_file)
            st.image(image, caption="Uploaded Prescription", use_column_width=True)
            
            if st.button("Extract Text" if language == "English" else "متن نکالیں"):
                with st.spinner("Processing prescription..." if language == "English" else "نسخہ پروسیس ہو رہا ہے..."):
                    extracted_text = ocr_processor.extract_text(image)
                    
                    if extracted_text:
                        if language == "English":
                            st.success("✅ Text extracted successfully!")
                            st.subheader("Extracted Text:")
                        else:
                            st.success("✅ متن کامیابی سے نکال لیا گیا!")
                            st.subheader("نکالا گیا متن:")
                        
                        st.text_area("", extracted_text, height=200)
                        
                        # Simple accuracy estimation (in real scenario, use validation dataset)
                        word_count = len(extracted_text.split())
                        non_empty_chars = len([c for c in extracted_text if c.strip()])
                        
                        if word_count > 5 and non_empty_chars > 20:
                            estimated_accuracy = min(85, (word_count / max(1, word_count)) * 100)
                            st.metric("Estimated OCR Accuracy", f"{estimated_accuracy:.1f}%")
                        else:
                            st.warning("Low confidence in OCR extraction")
                    
                    else:
                        if language == "English":
                            st.error("No text could be extracted from the image")
                        else:
                            st.error("تصویر سے کوئی متن نہیں نکالا جا سکا")
    
    with tab3:
        # Healthcare Chatbot
        if language == "English":
            st.header("Healthcare Assistant Chatbot")
            st.write("Ask health-related questions and get personalized advice")
        else:
            st.header("ہیلتھ کیئر اسسٹنٹ چیٹ بوٹ")
            st.write("صحت سے متعلق سوالات پوچھیں اور ذاتی مشورہ حاصل کریں")
        
        # Initialize chat history
        if 'chat_history' not in st.session_state:
            st.session_state.chat_history = []
        
        # Display chat history
        for message in st.session_state.chat_history:
            with st.chat_message(message["role"]):
                st.markdown(message["content"])
        
        # Chat input
        if prompt := st.chat_input("Type your health question here..." if language == "English" else "اپنا صحت کا سوال یہاں ٹائپ کریں..."):
            # Add user message to chat history
            st.session_state.chat_history.append({"role": "user", "content": prompt})
            with st.chat_message("user"):
                st.markdown(prompt)
            
            # Generate bot response
            with st.chat_message("assistant"):
                with st.spinner("Thinking..." if language == "English" else "سوچ رہا ہوں..."):
                    response = chatbot.get_response(prompt, language)
                    st.markdown(response)
            
            # Add assistant response to chat history
            st.session_state.chat_history.append({"role": "assistant", "content": response})
    
    with tab4:
        # Analytics Dashboard
        if language == "English":
            st.header("System Analytics & Performance")
        else:
            st.header("سسٹم تجزیات اور کارکردگی")
        
        col9, col10, col11 = st.columns(3)
        
        with col9:
            # Mock accuracy metrics (replace with actual validation)
            st.metric("Diagnostic Accuracy", "87%", "2%")
        with col10:
            st.metric("OCR Accuracy", "82%", "3%")
        with col11:
            st.metric("Risk Scoring AUC", "0.85", "0.02")
        
        # Performance charts
        if language == "English":
            st.subheader("Priority Distribution")
        else:
            st.subheader("ترجیحی تقسیم")
        
        # Mock data for demonstration
        priority_data = pd.DataFrame({
            'Priority': ['Emergency', 'Same Day', 'Routine'],
            'Count': [25, 45, 30]
        })
        
        fig = px.pie(priority_data, values='Count', names='Priority', 
                     title="Patient Priority Distribution")
        st.plotly_chart(fig, use_container_width=True)
        
        # Model performance
        if language == "English":
            st.subheader("Model Performance Metrics")
        else:
            st.subheader("ماڈل کارکردگی کے پیمانے")
        
        performance_data = pd.DataFrame({
            'Model': ['Heart Disease', 'Diabetes', 'Hypertension'],
            'Accuracy': [0.88, 0.85, 0.86],
            'AUC': [0.89, 0.84, 0.87]
        })
        
        fig = px.bar(performance_data, x='Model', y=['Accuracy', 'AUC'],
                     title="Model Performance Comparison", barmode='group')
        st.plotly_chart(fig, use_container_width=True)

if __name__ == "__main__":
    main()