import os
import gradio as gr
import streamlit as st
from groq import Groq
import numpy as np
from PIL import Image
from tensorflow.keras.models import load_model

# Load Pneumonia Detection Model
model = load_model('xray_image_classifier_model.keras')

# Set up Groq API Key
GROQ_API_KEY = "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"
os.environ["GROQ_API_KEY"] = GROQ_API_KEY

# Initialize the Groq client
client = Groq(api_key=GROQ_API_KEY)

# Define solutions
solutions = {
    "Pneumonia": "Consult a doctor immediately. Follow prescribed antibiotics if given, rest well, and stay hydrated.",
    "Normal": "Your X-ray appears normal. However, if you experience symptoms, consult a doctor for further evaluation."
}

# Prediction Function
def predict(image):
    img = image.resize((150, 150))
    img_array = np.array(img) / 255.0
    img_array = np.expand_dims(img_array, axis=0)
    
    prediction = model.predict(img_array)
    predicted_class = "Pneumonia" if prediction > 0.5 else "Normal"
    
    # Get the corresponding solution
    solution = solutions.get(predicted_class, "No specific advice available.")
    
    return predicted_class, solution

# CSS Styling for Gradio
css = """
    .gradio-container {
        background-color: #f5f5f5;
        font-family: Arial, sans-serif;
    }
    
    .gr-button {
        background-color:#007bff;
        color: white; 
        border: none;
        border-radius: 5px;
        font-size: 16px;
        padding: 10px 20px;
        cursor: pointer;
        transition: background-color 0.3s ease;
    }

    .gr-button:hover {
        background-color: #0056b3;
    }

    .gr-textbox, .gr-image {
        border: 2px dashed #007bff;
        padding: 20px;
        border-radius: 10px;
        background-color: #ffffff;
    }

    .gr-box-text {
        color: #007bff;
        font-size: 22px;
        font-weight: bold;
        text-align: center;
    }

    h1 {
        font-size: 36px;
        color: #007bff;
        text-align: center;
    }

    p {
        font-size: 20px;
        color: #333;
        text-align: center;
    }
"""

# Gradio UI for Pneumonia Detection
with gr.Blocks(css=css) as gradio_interface:
    gr.Markdown("<h1>Automated Pneumonia Detection via Chest X-ray Classification</h1>")
    gr.Markdown("<p>Submit a chest X-ray image below.</p>")
    
    with gr.Row():
        image_input = gr.Image(label="Drop Image Here", type="pil", elem_classes=["gr-image", "gr-box-text"])
        output_prediction = gr.Textbox(label="Model Analysis Output", elem_classes=["gr-textbox", "gr-box-text"])
        output_solution = gr.Textbox(label="Recommended Solution", elem_classes=["gr-textbox", "gr-box-text"])

    submit_btn = gr.Button("Initiate Diagnostic Analysis", elem_classes=["gr-button"])
    submit_btn.click(fn=predict, inputs=image_input, outputs=[output_prediction, output_solution])
    
    gr.Markdown("<h3>Note:</h3> <p>The AI model provides an initial assessment. Always consult a doctor for final diagnosis.</p>")

# Streamlit UI for Disease Chatbot
st.set_page_config(page_title="AI Health Assistant", page_icon="🩺", layout="wide")
st.title("🩺 AI Health Assistant")
st.write("Welcome! Upload an X-ray for pneumonia detection or ask the chatbot about diseases.")

# Sidebar Theme Settings
st.sidebar.header("⚙️ Settings")
chat_theme = st.sidebar.radio("Choose a theme:", ["Light", "Dark", "Blue", "Green"])

if chat_theme == "Dark":
    st.markdown("""
        <style>
            body {background-color: #1e1e1e; color: white;}
            .stButton>button {background-color: #4CAF50; color: white;}
            .chat-bubble {background-color: #2c2c2c; border-radius: 10px; padding: 10px;}
        </style>
    """, unsafe_allow_html=True)
elif chat_theme == "Blue":
    st.markdown("""
        <style>
            body {background-color: #e3f2fd; color: black;}
            .stButton>button {background-color: #2196F3; color: white;}
            .chat-bubble {background-color: #bbdefb; border-radius: 10px; padding: 10px;}
        </style>
    """, unsafe_allow_html=True)
elif chat_theme == "Green":
    st.markdown("""
        <style>
            body {background-color: #e8f5e9; color: black;}
            .stButton>button {background-color: #4CAF50; color: white;}
            .chat-bubble {background-color: #c8e6c9; border-radius: 10px; padding: 10px;}
        </style>
    """, unsafe_allow_html=True)

# Chatbot Function
def generate_chatbot_response(user_message):
    if "who created you" in user_message.lower():
        return "I was created by Abdel Basit. 😊"
    
    prompt = f"You are a helpful AI chatbot for medical guidance. The user is asking: {user_message}. Provide a detailed, professional response."
    
    chat_completion = client.chat.completions.create(
        messages=[{"role": "user", "content": prompt}],
        model="llama3-8b-8192",
    )
    
    return chat_completion.choices[0].message.content

# Chatbot Interface
st.markdown("### 💬 Chat with the AI Health Assistant")
user_input = st.chat_input("Ask me a health-related question:")

if user_input:
    chatbot_response = generate_chatbot_response(user_input)
    st.markdown(f"**You:** {user_input}")
    st.markdown(f"**AI:** {chatbot_response}")

# Launch Gradio Interface in Streamlit
st.markdown("---")
st.markdown("## 🔬 Pneumonia Detection System")
with st.expander("Click here to open the AI Pneumonia Detection System"):
    gradio_interface.launch()