import streamlit as st
import shutil 
import tensorflow as tf
from tensorflow.keras.models import load_model
import numpy as np
from PIL import Image
from huggingface_hub import hf_hub_download
import os

HF_REPO_ID = "saad1BM/brain-tumor-detection" 
MODEL_FILENAME = "brain_tumor_model.keras" 
MODEL_PATH = MODEL_FILENAME 

IMAGE_SIZE = (224, 224)
CLASS_NAMES = ['glioma_tumor', 'meningioma_tumor', 'no_tumor', 'pituitary_tumor']


@st.cache_resource
def load_trained_model():
    """Trained model ko Hugging Face se download aur load karta hai."""
    
   
    global HF_REPO_ID, MODEL_FILENAME, MODEL_PATH
    
    MIN_SIZE = 100 * 1024 * 1024 

    try:
        
        if not os.path.exists(MODEL_PATH) or os.path.getsize(MODEL_PATH) < MIN_SIZE: 
            st.info(f"Downloading model from Hugging Face: {HF_REPO_ID}/{MODEL_FILENAME}")
            
          
            downloaded_path = hf_hub_download(
                repo_id=HF_REPO_ID, 
                filename=MODEL_FILENAME, 
                
                cache_dir="/tmp/hf_cache", 
            )
            
            st.info("Copying model file locally to avoid cross-device link error...")
            
            
            shutil.copyfile(downloaded_path, MODEL_PATH) 

            st.success("Model download and copy successful!")
            
        
        model = load_model(MODEL_PATH)
        return model
        
    except Exception as e:
        
        st.error(f"Model Load Failed. Ensure 'huggingface-hub' is in requirements.txt and file is public. Error: {e}")
        return None

model = load_trained_model()


def predict_image(image_file, model):
    """Uploaded image par prediction karta hai."""
    if model is None:
        return "Model Load Failed", 0.0

    
    img = Image.open(image_file).convert("RGB")
    
    img = img.resize(IMAGE_SIZE)
    img_array = np.array(img)
    
    
    img_array = np.expand_dims(img_array, axis=0)
    
    
    img_array = img_array / 255.0

    
    predictions = model.predict(img_array)
    
    
    predicted_index = np.argmax(predictions, axis=1)[0]
    confidence_score = np.max(predictions) * 100

    predicted_class = CLASS_NAMES[predicted_index]
    
    return predicted_class, confidence_score



st.set_page_config(page_title="Brain Tumor Detection", layout="wide")


st.title("🧠 Brain Tumor Detection System (AI Powered)")
st.write("Upload an MRI image below to classify it as one of the tumor types or no tumor.")
st.markdown("---")

col1, col2 = st.columns(2)

with col1:
    
    uploaded_file = st.file_uploader("Upload MRI Image:", type=["jpg", "jpeg", "png"])
    
    if uploaded_file is not None:
        
        st.image(uploaded_file, caption="Uploaded MRI Image", use_column_width=True)
        st.markdown("---")

        
        if st.button("Detect Tumor"):
            st.spinner("Analyzing image and detecting tumor...")
            
            
            predicted_class, confidence_score = predict_image(uploaded_file, model)
            
            
            
            
            if predicted_class == 'no_tumor':
                result_label = f"🟢 **Prediction: No Tumor**"
            elif predicted_class == 'Model Load Failed':
                result_label = f"❌ **Prediction: Model Initialization Error**"
                st.error("Model could not be loaded for prediction. Please check logs for Hugging Face download errors.")
            else:
                result_label = f"🔴 **Prediction: Tumor ({predicted_class.replace('_', ' ').title()})**"
            
            st.success("✅ Analysis Complete")
            st.subheader(result_label)
            st.metric(label="Confidence Score", value=f"{confidence_score:.2f}%")
            
            st.write("---")

with col2:
    st.header("Results and Interpretation")
    st.info("The system uses Transfer Learning (VGG16) to classify the image into four categories: Glioma, Meningioma, Pituitary, or No Tumor.")
    
    if uploaded_file is None:
        st.warning("Please upload an image and click 'Detect Tumor' to see the results.")