import streamlit as st
import numpy as np
import cv2
import tensorflow as tf
from PIL import Image
from keras import backend as K
from keras.metrics import Precision, Recall
from vit_keras import vit, utils, layers

# Page configuration
st.set_page_config(
    page_title="Breast Cancer Classification",
    page_icon="🏥",
    layout="centered"
)

# Cache the model loading
@st.cache_resource
def load_model():
    try:
        model = tf.keras.models.load_model(
            'model.h5',
            custom_objects={
                'ClassToken': layers.ClassToken,
                'f1_score': f1_score
            },
            compile=False
        )
        return model
    except Exception as e:
        st.error(f"Error loading model: {str(e)}")
        return None

def f1_score(y_true, y_pred):
    y_true = K.round(y_true)
    y_pred = K.round(y_pred)
    
    tp = K.sum(y_true * y_pred)
    fp = K.sum((1 - y_true) * y_pred)
    fn = K.sum(y_true * (1 - y_pred))
    
    precision = tp / (tp + fp + K.epsilon())
    recall = tp / (tp + fn + K.epsilon())
    
    return 2 * precision * recall / (precision + recall + K.epsilon())

def process_image(image):
    img_array = np.array(image)
    image_pred = cv2.cvtColor(img_array, cv2.COLOR_BGR2RGB)
    resized_image = cv2.resize(image_pred, (224, 224))
    return np.array([resized_image], dtype='float32') / 255.0

# App header
st.header('🔬 Breast Cancer Classification using Ultrasound')

# Load model
model = load_model()

if model:
    model.compile(
        optimizer='adam',
        loss='binary_crossentropy',
        metrics=['accuracy', f1_score]
    )

    # File uploader
    uploaded_file = st.file_uploader(
        "Upload an ultrasound image",
        type=['jpg', 'jpeg', 'png'],
        help="Supported formats: JPG, JPEG, PNG"
    )

    if uploaded_file:
        try:
            # Display image
            col1, col2 = st.columns(2)
            with col1:
                image = Image.open(uploaded_file)
                st.image(image, caption='Uploaded Ultrasound Image', use_column_width=True)
            
            # Process and predict
            with st.spinner('Analyzing image...'):
                processed_image = process_image(image)
                predictions = model.predict(processed_image, verbose=0)
                predicted_label_index = tf.argmax(predictions, axis=-1).numpy()[0]
            
            # Display results
            with col2:
                st.subheader("Analysis Results")
                labels = {'Benign': 0, 'Malignant': 1}
                result = list(labels.keys())[list(labels.values()).index(predicted_label_index)]
                confidence = float(predictions[0][predicted_label_index]) * 100
                
                st.metric("Prediction", result)
                st.metric("Confidence", f"{confidence:.2f}%")
                
                if predicted_label_index == 1:
                    st.error("⚠️ Please consult with your doctor as results indicate potential malignancy.")
                else:
                    st.success("✅ Results suggest benign characteristics.")

        except Exception as e:
            st.error(f"Error processing image: {str(e)}")
else:
    st.error("Failed to load the model. Please check if model.h5 exists in the correct location.")