import gradio as gr
import numpy as np
import pickle
import time

# Load the trained models
with open("stroke_clf.pkl", "rb") as f:
    clf_model = pickle.load(f)

with open("stroke_reg.pkl", "rb") as f:
    reg_model = pickle.load(f)

# Custom CSS for enhanced styling
custom_css = """
.gradio-container {
    max-width: 1200px !important;
    margin-left: auto !important;
    margin-right: auto !important;
}

.container {
    margin: 0 auto !important;
    padding: 2rem !important;
}

.question-group {
    border: 1px solid #e5e7eb !important;
    border-radius: 8px !important;
    padding: 1.5rem !important;
    margin-bottom: 1.5rem !important;
    background: white !important;
}

.footer {
    text-align: center !important;
    padding: 2rem !important;
    background: #f8fafc !important;
    border-top: 1px solid #e5e7eb !important;
}
"""

def predict_stroke(*inputs):
    """
    Enhanced prediction function with detailed output formatting
    """
    age = inputs[0]
    features = [1 if x == "Yes" else 0 for x in inputs[1:]] + [age]
    
    sample_input = np.array([features])
    
    # Get predictions
    classification_result = clf_model.predict(sample_input)[0]
    regression_result = reg_model.predict(sample_input)[0]
    risk_percentage = round(regression_result, 2)
    
    # Simulate processing
    time.sleep(1.5)
    
    # Enhanced results formatting
    if classification_result == 1:
        severity = "High" if risk_percentage > 70 else "Moderate"
        color = "red" if risk_percentage > 70 else "orange"
        result = f"""
        <div style='background-color: #FEF2F2; padding: 20px; border-radius: 10px; border: 1px solid {color}'>
            <h2 style='color: {color}; margin-bottom: 15px'>🚨 Risk Assessment Results: {severity} Risk</h2>
            <div style='background-color: white; padding: 15px; border-radius: 8px; margin-bottom: 15px'>
                <p style='font-size: 18px; margin-bottom: 10px'><strong>Risk Level:</strong> {risk_percentage}%</p>
                <p style='color: {color}'><strong>Status:</strong> Immediate Attention Recommended</p>
            </div>
            
            <h3 style='margin: 20px 0 10px 0'>📋 Recommended Actions:</h3>
            <ol style='margin-left: 20px'>
                <li>Schedule an urgent appointment with your healthcare provider</li>
                <li>Begin daily blood pressure monitoring</li>
                <li>Review and log all current medications</li>
                <li>Start a symptom diary</li>
            </ol>

            <h3 style='margin: 20px 0 10px 0'>🌟 Immediate Lifestyle Changes:</h3>
            <ul style='margin-left: 20px'>
                <li>Reduce sodium intake</li>
                <li>Increase water consumption</li>
                <li>Practice stress-reduction techniques</li>
                <li>Ensure 7-8 hours of quality sleep</li>
            </ul>
        </div>
        """
    else:
        result = f"""
        <div style='background-color: #F0FDF4; padding: 20px; border-radius: 10px; border: 1px solid #22C55E'>
            <h2 style='color: #22C55E; margin-bottom: 15px'>✅ Risk Assessment Results: Low Risk</h2>
            <div style='background-color: white; padding: 15px; border-radius: 8px; margin-bottom: 15px'>
                <p style='font-size: 18px; margin-bottom: 10px'><strong>Risk Level:</strong> {risk_percentage}%</p>
                <p style='color: #22C55E'><strong>Status:</strong> Healthy Range</p>
            </div>
            
            <h3 style='margin: 20px 0 10px 0'>🌟 Maintaining Your Health:</h3>
            <ul style='margin-left: 20px'>
                <li>Continue regular exercise routine</li>
                <li>Maintain balanced nutrition</li>
                <li>Schedule regular check-ups</li>
                <li>Practice stress management</li>
            </ul>

            <h3 style='margin: 20px 0 10px 0'>💪 Preventive Measures:</h3>
            <ul style='margin-left: 20px'>
                <li>Annual health screenings</li>
                <li>Regular blood pressure checks</li>
                <li>Adequate sleep and hydration</li>
                <li>Active lifestyle maintenance</li>
            </ul>
        </div>
        """
    
    return result

# Create the enhanced Gradio interface
with gr.Blocks(
    theme=gr.themes.Soft(primary_hue="blue", secondary_hue="purple"),
    title="EarlyMed: AI-Powered Stroke Risk Assessment",
    css=custom_css
) as demo:
    
    # Header Section
    gr.HTML("""
        <div style="text-align: center; padding: 2rem; background: linear-gradient(to right, #E8F0FF, #F0F7FF);">
            <h1 style="color: #1E40AF; margin-bottom: 1rem;">🌟 EarlyMed: AI-Powered Stroke Risk Assessment 🧠</h1>
            <p style="color: #4B5563; font-size: 1.1rem; max-width: 800px; margin: 0 auto;">
                Welcome to <strong>EarlyMed</strong>—a VIT-AP University initiative using advanced AI to help you understand your stroke risk factors.
            </p>
        </div>
    """)
    
    with gr.Row():
        with gr.Column():
            # Basic Information
            with gr.Group(elem_classes="question-group"):
                gr.Markdown("### 👤 Your Profile")
                age = gr.Slider(
                    label="How old are you?",
                    minimum=18,
                    maximum=100,
                    step=1,
                    value=50,
                    info="Move the slider to your age"
                )

            # Primary Symptoms
            with gr.Group(elem_classes="question-group"):
                gr.Markdown("### 💗 Heart & Chest Symptoms")
                chest_pain = gr.Radio(
                    label="Do you experience any discomfort or pain in your chest?",
                    choices=["No", "Yes"],
                    value="No",
                    info="This might feel like pressure, squeezing, or general discomfort"
                )
                shortness_breath = gr.Radio(
                    label="Do you find yourself getting out of breath easily?",
                    choices=["No", "Yes"],
                    value="No",
                    info="For example, during light activities like climbing stairs"
                )
                irregular_heartbeat = gr.Radio(
                    label="Have you noticed your heart beating irregularly?",
                    choices=["No", "Yes"],
                    value="No",
                    info="This might feel like skipped beats or fluttering"
                )

        with gr.Column():
            # Daily Symptoms
            with gr.Group(elem_classes="question-group"):
                gr.Markdown("### 😮‍💨 Daily Experiences")
                fatigue_weakness = gr.Radio(
                    label="Do you often feel unusually tired or weak?",
                    choices=["No", "Yes"],
                    value="No"
                )
                dizziness = gr.Radio(
                    label="Do you experience dizziness or light-headedness?",
                    choices=["No", "Yes"],
                    value="No"
                )
                swelling = gr.Radio(
                    label="Have you noticed swelling in your legs or ankles?",
                    choices=["No", "Yes"],
                    value="No"
                )

    with gr.Row():
        with gr.Column():
            # Additional Symptoms
            with gr.Group(elem_classes="question-group"):
                gr.Markdown("### 🤒 Additional Signs")
                pain_neck_jaw = gr.Radio(
                    label="Do you experience pain in your neck, jaw, or shoulders?",
                    choices=["No", "Yes"],
                    value="No"
                )
                excessive_sweating = gr.Radio(
                    label="Do you sweat more than usual, even when not active?",
                    choices=["No", "Yes"],
                    value="No"
                )
                persistent_cough = gr.Radio(
                    label="Have you developed a cough that won't go away?",
                    choices=["No", "Yes"],
                    value="No"
                )

        with gr.Column():
            # Health Conditions
            with gr.Group(elem_classes="question-group"):
                gr.Markdown("### 🏥 Health Conditions")
                nausea_vomiting = gr.Radio(
                    label="Do you feel nauseated or sick to your stomach?",
                    choices=["No", "Yes"],
                    value="No"
                )
                high_bp = gr.Radio(
                    label="Has a doctor ever told you that you have high blood pressure?",
                    choices=["No", "Yes"],
                    value="No"
                )
                chest_discomfort = gr.Radio(
                    label="Do you feel discomfort in your chest when you're active?",
                    choices=["No", "Yes"],
                    value="No"
                )

    with gr.Row():
        with gr.Column():
            # Lifestyle Factors
            with gr.Group(elem_classes="question-group"):
                gr.Markdown("### 🌟 Lifestyle & Well-being")
                cold_hands_feet = gr.Radio(
                    label="Do your hands or feet often feel unusually cold?",
                    choices=["No", "Yes"],
                    value="No"
                )
                snoring = gr.Radio(
                    label="Has anyone told you that you snore loudly or stop breathing during sleep?",
                    choices=["No", "Yes"],
                    value="No"
                )
                anxiety = gr.Radio(
                    label="Do you often feel anxious or experience a sense of doom?",
                    choices=["No", "Yes"],
                    value="No"
                )

    # Submit Button
    gr.Markdown("""<div style='text-align: center; margin: 2rem 0;'>""")
    submit_btn = gr.Button(
        "Analyze My Risk Factors",
        variant="primary",
        size="lg"
    )
    
    # Output
    output = gr.HTML(label="Assessment Results")
    
    # Wire up the prediction
    submit_btn.click(
        predict_stroke,
        inputs=[age, chest_pain, shortness_breath, irregular_heartbeat,
                fatigue_weakness, dizziness, swelling, pain_neck_jaw,
                excessive_sweating, persistent_cough, nausea_vomiting,
                high_bp, chest_discomfort, cold_hands_feet, snoring, anxiety],
        outputs=output
    )
    
    # Footer
    gr.HTML("""
        <div style="text-align: center; margin-top: 2rem; padding: 2rem; background: linear-gradient(to right, #F8FAFC, #F0F7FF);">
            <h3 style="color: #1E40AF; margin-bottom: 1rem;">💡 Important Note</h3>
            <p style="color: #4B5563; max-width: 800px; margin: 0 auto 1rem auto;">
                This tool is designed for educational purposes and early awareness only. It is not a substitute for professional medical diagnosis or advice.
            </p>
            <p style="color: #6B7280; font-size: 0.9rem;">
                © 2025 EarlyMed - VIT-AP University | All Rights Reserved
            </p>
        </div>
    """)

# Launch the interface
demo.launch()