import streamlit as st
import numpy as np
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing import image
import matplotlib.pyplot as plt

# Disable the PyplotGlobalUseWarning
st.set_option('deprecation.showPyplotGlobalUse', False)

# Load the model
model = load_model('brain_tumor_model.keras')

# Class mappings
class_mappings = {
    'Glioma': 0,
    'Meningioma': 1,
    'Notumor': 2,
    'Pituitary': 3
}
inv_class_mappings = {v: k for k, v in class_mappings.items()}
class_names = list(class_mappings.keys())

# Load the true and predicted labels
true_labels = np.load('true_labels.npy') # Model trained on Brain Tumor MRI Dataset
predicted_labels = np.load('predicted_labels.npy') # Model foments predictions based upon the training/testing of the Brain Tumor MRI Dataset

# Note0: The One Hot Encode Predictions are leveraging the TF CNN Model's training/testing of the Brain Tumor MRI Dataset
# Note1: The One Hot Encode Predictions are deriving from another dataset -> Crystal Clean: Brain Tumors MRI Dataset


# Function to load and preprocess an image
def load_and_preprocess_image(image_path, image_shape=(168, 168)):
    img = image.load_img(image_path, target_size=image_shape, color_mode='grayscale')
    img_array = image.img_to_array(img) / 255.0
    img_array = np.expand_dims(img_array, axis=0)
    return img_array

# Function to display a row of images with predictions
def display_images_and_predictions(image_paths, predictions, true_labels, figsize=(20, 5)):
    fig, axes = plt.subplots(1, len(image_paths), figsize=figsize)
    for i, (image_path, prediction, true_label) in enumerate(zip(image_paths, predictions, true_labels)):
        ax = axes[i]
        img_array = load_and_preprocess_image(image_path)
        img_array = np.squeeze(img_array)
        ax.imshow(img_array, cmap='gray')
        title_color = 'green' if prediction == true_label else 'red'
        ax.set_title(f'True Label: {true_label}\nPred: {prediction}', color=title_color)
        ax.axis('off')
    st.pyplot(fig)



# Image paths
image_paths = [
    'N_1.jpg',
    'G_1.jpg',
    'M_1.jpg',
    'P_1.jpg',
    #'N_1_BR_.jpg',
    #'G_1_BR_.jpg',
    #'M_1_BR_.jpg',
    #'P_1_BR_.jpg',
   #'N_1_DA_.jpg',
    #'G_1_DA_.jpg',
    #'M_1_DA_.jpg',
    #'P_1_DA_.jpg',
    #'N_1_HF_.jpg',
    #'G_1_HF_.jpg',
    #'M_1_HF_.jpg',
    #'P_1_HF_.jpg',
    #'N_1_RO_.jpg',
    #'G_1_RO_.jpg',
    #'M_1_RO_.jpg',
    #'P_1_RO_.jpg',
    #'N_1_SP_.jpg',
    #'G_1_SP_.jpg',
    #'M_1_SP_.jpg',
    #'P_1_SP_.jpg',
    #'N_1_VF_.jpg',
    #'G_1_VF_.jpg',
    #'M_1_VF_.jpg',
    #'P_1_VF_.jpg'
   
]

# True labels for images
true_labels = [
    'Notumor', 'Glioma', 'Meningioma', 'Pituitary',  # Original tumor types
    #'Notumor', 'Glioma', 'Meningioma', 'Pituitary',  # BR_ (Brightness)
    #'Notumor', 'Glioma', 'Meningioma', 'Pituitary',  # DA_ (Darkness)
    #'Notumor', 'Glioma', 'Meningioma', 'Pituitary',  # HF_ (High Frequency)
    #'Notumor', 'Glioma', 'Meningioma', 'Pituitary',  # RO_ (Rotation)
    #'Notumor', 'Glioma', 'Meningioma', 'Pituitary',  # SP_ (Salt and Pepper)
    #'Notumor', 'Glioma', 'Meningioma', 'Pituitary'   # VF_ (Vertical Flip)
]

# Load and preprocess images, then make predictions
images = [load_and_preprocess_image(path) for path in image_paths]
predictions = [model.predict(image) for image in images]

# Determine the predicted labels
predicted_labels = [inv_class_mappings[np.argmax(one_hot)] for one_hot in predictions]

# Create Streamlit app title
st.markdown("<h1 style='text-align: center; color: navy;'>Brain Tumor One Hot Encode TF Model</h1>", unsafe_allow_html=True)

# Output the predictions
#st.write(f'Class Mappings: {class_mappings}')
#for i, (prediction, true_label) in enumerate(zip(predicted_labels, true_labels)):
    #st.write(f"Image {i+1} Prediction: {prediction}, True Label: {true_label}")

# Output the predictions
st.write(f'Class Mappings: {class_mappings}')
st.write("\nNormal Image Prediction:", np.round(predictions[0], 3)[0])
st.write("Glioma Image Prediction:", np.round(predictions[1], 3)[0])
st.write("Meningioma Image Prediction:", np.round(predictions[2], 3)[0])
st.write("Pituitary Image Prediction:", np.round(predictions[3], 3)[0])


# Display images with predictions
display_images_and_predictions(image_paths, predicted_labels, true_labels)