import gradio as gr
import pandas as pd
import torch
from predict import RiskPredictor

# Initialize the predictor
predictor = RiskPredictor()

def predict_risk(age, bmi, systolic_bp, diastolic_bp, cholesterol, heart_rate, 
                 smoking, steps, stress, physical_activity, sleep, family_history, 
                 diet_quality, alcohol, risk_score):
    
    input_data = {
        'age': age,
        'bmi': bmi,
        'systolic_bp': systolic_bp,
        'diastolic_bp': diastolic_bp,
        'cholesterol_mg_dl': cholesterol,
        'resting_heart_rate': heart_rate,
        'smoking_status': smoking,
        'daily_steps': steps,
        'stress_level': stress,
        'physical_activity_hours_per_week': physical_activity,
        'sleep_hours': sleep,
        'family_history_heart_disease': family_history,
        'diet_quality_score': diet_quality,
        'alcohol_units_per_week': alcohol,
        'heart_disease_risk_score': risk_score
    }
    
    prediction = predictor.predict_single(input_data)
    
    # Return mapping for color-coded feedback
    result_text = f"## Predicted Category: {prediction.upper()}"
    
    if prediction == 'Low':
        description = "✅ Low risk! Excellent heart health habits."
        color = "#2ed573"
    elif prediction == 'Medium':
        description = "⚠️ Moderate risk. Consider heart-healthy changes."
        color = "#ffa502"
    else:
        description = "🚨 High risk! Please consult a health professional."
        color = "#ff4757"
        
    return f'<div style="background-color: {color}; padding: 20px; border-radius: 10px; color: white;">{result_text}<br>{description}</div>'

# Gradio Interface
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🏥 CardioGuard: RNN Heart Risk Predictor")
    gr.Markdown("Enter patient data below to analyze cardiovascular risk using a Deep Learning RNN model.")
    
    with gr.Row():
        with gr.Column():
            age = gr.Number(label="Age", value=45)
            bmi = gr.Number(label="BMI", value=26.5)
            systolic = gr.Number(label="Systolic BP", value=130)
            diastolic = gr.Number(label="Diastolic BP", value=85)
            cholesterol = gr.Number(label="Cholesterol (mg/dl)", value=210)
        
        with gr.Column():
            smoking = gr.Dropdown(["Never", "Former", "Current"], label="Smoking Status", value="Never")
            family_history = gr.Dropdown(["No", "Yes"], label="Family History", value="No")
            steps = gr.Number(label="Daily Steps", value=6000)
            heart_rate = gr.Number(label="Resting Heart Rate", value=72)
            risk_score = gr.Number(label="Internal Risk Score (0-100)", value=35.0)

        with gr.Column():
            stress = gr.Slider(1, 10, step=1, label="Stress Level", value=5)
            diet = gr.Slider(1, 10, step=1, label="Diet Quality", value=6)
            activity = gr.Number(label="Physical Activity (hrs/wk)", value=3.5)
            sleep = gr.Number(label="Sleep Hours", value=7.5)
            alcohol = gr.Number(label="Alcohol Units/wk", value=2.0)

    btn = gr.Button("Analyze Risk Profile", variant="primary")
    output = gr.HTML()

    btn.click(predict_risk, inputs=[
        age, bmi, systolic, diastolic, cholesterol, heart_rate, 
        smoking, steps, stress, activity, sleep, family_history, 
        diet, alcohol, risk_score
    ], outputs=output)

if __name__ == "__main__":
    demo.launch()