import streamlit as st
import tensorflow as tf
import random
from PIL import Image
from tensorflow import keras
import numpy as np
import os

import warnings

warnings.filterwarnings("ignore")
os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'

st.set_page_config(
    page_title="ChestAI - Pneumonia Detection",
    page_icon="🫁",
    initial_sidebar_state="auto",
)

hide_streamlit_style = """
	<style>
  #MainMenu {visibility: hidden;}
	footer {visibility: hidden;}
  </style>
"""
st.markdown(hide_streamlit_style, unsafe_allow_html=True)


def prediction_cls(prediction):
    for key, clss in class_names.items():
        if np.argmax(prediction) == clss:
            return key


with st.sidebar:
    st.title("ChestAI")
    st.markdown("""
        ### About
        ChestAI uses advanced deep learning to detect pneumonia in chest X-rays.
                
        ### How to use
        1. Upload a chest X-ray image (JPG/PNG)
        2. Wait for the analysis
        3. View the results and confidence score
                
        ### Note
        This tool is for educational purposes only. Always consult healthcare professionals for medical advice.
    """)
st.set_option("deprecation.showfileUploaderEncoding", False)


@st.cache_resource(show_spinner=False)
def load_model():
    try:
        from huggingface_hub import from_pretrained_keras
        config = tf.compat.v1.ConfigProto()
        config.gpu_options.allow_growth = True
        tf.compat.v1.keras.backend.set_session(tf.compat.v1.Session(config=config))
        
        keras.utils.set_random_seed(42)
        model = from_pretrained_keras("ryefoxlime/PneumoniaDetection")
        return model
    except Exception as e:
        st.error(f"Error loading model: {str(e)}")
        return None


with st.spinner("Model is being loaded..."):
    model = load_model()

if model is None:
    st.error("Failed to load model. Please try again.")
    st.stop()

file = st.file_uploader(" ", type=["jpg", "png"])


def import_and_predict(image_data, model):
    img_array = keras.preprocessing.image.img_to_array(image_data)
    img_array = np.expand_dims(img_array, axis=0)
    img_array = img_array/255

    predictions = model.predict(img_array)
    return predictions


if file is None:
    st.text("Please upload an image file")
else:
    try:
        image = keras.preprocessing.image.load_img(file, target_size=(224, 224), color_mode='rgb')
        st.image(image, caption="Uploaded Image.", use_column_width=True)
        predictions = import_and_predict(image, model)
        
        class_names = [
            "Normal",
            "PNEUMONIA",
        ]

        confidence = float(max(predictions[0]) * 100)
        prediction_label = class_names[np.argmax(predictions)]
        
        st.info(f"Confidence: {confidence:.2f}%")
        
        if prediction_label == "Normal":
            st.balloons()
            st.success(f"Result: {prediction_label}")
        else:
            st.warning(f"Result: {prediction_label}")
            
    except Exception as e:
        st.error(f"Error processing image: {str(e)}")