# app.py
import streamlit as st
import numpy as np
import os
import tensorflow as tf
import logging
from PIL import Image

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Set page configuration
st.set_page_config(
    page_title="Breast Cancer Prediction",
    page_icon="🩺",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Disable GPU to save memory
tf.config.set_visible_devices([], 'GPU')
logger.info("TensorFlow configured for CPU-only")

# ===== Model Loading =====
MODEL_FILE = "final_combined_model.keras"

@st.cache_resource(show_spinner=False)
def load_model():
    """Load TensorFlow model from local file with caching"""
    try:
        # Verify file exists
        if not os.path.exists(MODEL_FILE):
            logger.error(f"❌ Model file not found: {MODEL_FILE}")
            return None
        
        logger.info(f"⏳ Loading model from local file: {MODEL_FILE}")
        
        # Load model with memory optimization
        model = tf.keras.models.load_model(MODEL_FILE, compile=False)
        
        # Test prediction to verify loading
        test_input = np.random.rand(1, 224, 224, 1).astype(np.float32)
        test_pred = model.predict(test_input, verbose=0)
        logger.info(f"🧪 Test prediction: {test_pred[0][0]:.4f}")
        
        logger.info("✅ Model loaded successfully")
        return model
    except Exception as e:
        logger.error(f"❌ Error loading model: {e}")
        # Print detailed traceback
        import traceback
        logger.error(traceback.format_exc())
        return None

# Load model at startup
model = load_model()

# ===== Image Preprocessing =====
def preprocess_image(image):
    """Preprocess image for model prediction"""
    try:
        # Convert to PIL Image
        if isinstance(image, np.ndarray):
            img = Image.fromarray(image.astype('uint8'))
        else:
            img = image
        
        # Processing pipeline
        img = img.convert('L')        # Grayscale
        img = img.resize((224, 224))  # Resize
        img_array = np.array(img) / 255.0  # Normalize
        
        # Add batch and channel dimensions
        return img_array[np.newaxis, ..., np.newaxis]
    except Exception as e:
        logger.error(f"🖼️ Image preprocessing error: {e}")
        return None

# ===== Prediction Function =====
def predict(image):
    """Make prediction using the loaded model"""
    if model is None:
        return "Model failed to load", "Check logs", None
    
    try:
        # Preprocess image
        processed_image = preprocess_image(image)
        if processed_image is None:
            return "Invalid image", "Try another", image
        
        # Make prediction
        prediction = model.predict(processed_image, verbose=0)[0][0]
        
        # Format results
        confidence = abs(prediction - 0.5) + 0.5  # Convert to 0.5-1.0 scale
        result = "Malignant" if prediction > 0.5 else "Benign"
        
        return result, f"{confidence*100:.2f}%", image
    except Exception as e:
        error_msg = f"Prediction error: {str(e)}"
        logger.error(error_msg)
        return error_msg, "Try again", image

# ===== Streamlit UI =====

# Custom CSS for styling
st.markdown("""
<style>
    .stApp {
        background-color: #f0f2f6;
    }
    .header {
        color: #2c3e50;
        text-align: center;
        padding: 1rem;
    }
    .result-box {
        border-radius: 10px;
        padding: 1.5rem;
        margin: 1rem 0;
        box-shadow: 0 4px 6px rgba(0,0,0,0.1);
    }
    .malignant {
        background-color: #ffcccc;
        border-left: 5px solid #e74c3c;
    }
    .benign {
        background-color: #ccffcc;
        border-left: 5px solid #2ecc71;
    }
    .stButton>button {
        background-color: #3498db;
        color: white;
        border-radius: 5px;
        padding: 0.5rem 1rem;
        width: 100%;
    }
    .stButton>button:hover {
        background-color: #2980b9;
    }
</style>
""", unsafe_allow_html=True)

# Header
st.markdown("<h1 class='header'>🩺 Breast Cancer Prediction</h1>", unsafe_allow_html=True)
st.markdown("Upload a breast medical image for cancer prediction")

# Status indicator
status = "✅ Model loaded successfully" if model else "❌ Model failed to load"
st.info(status)

# Create two columns for layout
col1, col2 = st.columns([1, 1])

# Input column
with col1:
    st.subheader("Patient Information")
    
    # Input fields
    age = st.number_input("Patient Age", min_value=18, max_value=100, value=45)
    tumor_size = st.number_input("Tumor Size (mm)", min_value=0.1, value=15.0)
    
    # Image upload
    uploaded_file = st.file_uploader(
        "Upload Medical Image", 
        type=["jpg", "jpeg", "png"],
        help="Supported formats: JPG, JPEG, PNG"
    )
    
    # Predict button
    predict_btn = st.button("Analyze Image")

# Results column
with col2:
    st.subheader("Prediction Results")
    
    # Initialize session state for results
    if 'result' not in st.session_state:
        st.session_state.result = None
        st.session_state.confidence = None
        st.session_state.image = None
    
    # Process image when button is clicked
    if predict_btn and uploaded_file is not None:
        try:
            image = Image.open(uploaded_file)
            st.session_state.result, st.session_state.confidence, st.session_state.image = predict(image)
        except Exception as e:
            st.error(f"Error processing image: {str(e)}")
    
    # Display results if available
    if st.session_state.result:
        # Result box with color coding
        result_class = "malignant" if st.session_state.result == "Malignant" else "benign"
        st.markdown(
            f"<div class='result-box {result_class}'>"
            f"<h3>Diagnosis: {st.session_state.result}</h3>"
            f"<p>Confidence: {st.session_state.confidence}</p>"
            "</div>", 
            unsafe_allow_html=True
        )
        
        # Display image
        if st.session_state.image:
            st.image(
                st.session_state.image, 
                caption="Uploaded Image",
                use_container_width=True
            )
    
    # Show placeholder if no results
    elif not predict_btn:
        st.info("Upload an image and click 'Analyze Image' to get prediction")

# Footer
st.markdown("---")
st.caption("This tool is for research purposes only. Consult a medical professional for clinical diagnosis.")