import logging
import streamlit as st
from tensorflow.keras.models import load_model
from image_preprocess import preprocess
import os
import numpy as np
import joblib
import matplotlib.pyplot as plt

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Load models
logger.info("Loading feature extractor model...")
feature_extractor = load_model("custom_cnn_feature_extractor_finale.h5")
logger.info("Feature extractor model loaded successfully.")

logger.info("Loading classification model...")
classification_model = joblib.load("RF_model_finale.pkl")
logger.info("Classification model loaded successfully.")

# Streamlit App
def main():
    st.title("Lung Cancer Classification Model")

    # Upload Image
    imagefile = st.file_uploader("Upload an Image", type=["jpg", "png", "jpeg"])

    if imagefile is not None:
        # Save the image temporarily
        image_path = os.path.join("static", "images", imagefile.name)
        with open(image_path, "wb") as f:
            f.write(imagefile.getbuffer())
        st.image(image_path, caption="Uploaded Image", use_container_width=True)

        try:
            # Preprocessing
            st.write("Starting image preprocessing...")
            X = preprocess(image_path)
            logger.info("Image preprocessing completed.")

            encode_label = {0: "Benign", 1: "Malignant", 2: "Normal"}

            # Extract Features
            logger.info("Extracting features using the feature extractor model...")
            features = feature_extractor.predict(X)
            features = features.reshape(features.shape[0], -1)
            logger.info("Feature extraction completed.")

            # Predict
            logger.info("Making predictions using the classification model...")
            Y_pred = classification_model.predict(features)
            probability = np.max(classification_model.predict_proba(features), axis=1) * 100
            Y_pred_index = int(Y_pred[0])
            probability_value = float(probability[0])
            classification = '%s (%.2f%%)' % (encode_label[Y_pred_index], probability_value)
            st.write(f"Prediction: {classification}")

            # Plot confidence chart
            st.write("Generating confidence chart...")
            probabilities = classification_model.predict_proba(features)[0] * 100
            classes = list(encode_label.values())

            plt.figure(figsize=(12, 6))
            bars = plt.bar(classes, probabilities, color=['blue', 'red', 'green'])
            for bar in bars:
                yval = bar.get_height()
                plt.text(bar.get_x() + bar.get_width() / 2, yval + 2, f'{yval:.1f}%', ha='center', va='bottom')
            plt.ylim(0, max(probabilities) + 10)
            plt.xlabel('Classes')
            plt.ylabel('Confidence Level (%)')
            plt.title('Confidence Level for Each Class')

            # Display chart in Streamlit
            st.pyplot(plt)

        except Exception as e:
            logger.error(f"Error during prediction: {str(e)}", exc_info=True)
            st.error(f"Error: {str(e)}")

if __name__ == "__main__":
    main()