"""
🐱🐶 Cat vs Dog Classifier - Streamlit Web App
================================================
Upload an image and get instant prediction
"""

import streamlit as st
import tensorflow as tf
from tensorflow.keras.preprocessing import image
from tensorflow.keras.models import load_model
import numpy as np
from PIL import Image
import os
import time

# ═══════════════════════════════════════════════════════════
# PAGE CONFIGURATION
# ═══════════════════════════════════════════════════════════
st.set_page_config(
    page_title="🐱 Cat vs Dog Classifier",
    page_icon="🐾",
    layout="centered"
)

# ═══════════════════════════════════════════════════════════
# CUSTOM CSS
# ═══════════════════════════════════════════════════════════
st.markdown("""
<style>
    .main-title {
        font-size: 3em;
        text-align: center;
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        -webkit-background-clip: text;
        -webkit-text-fill-color: transparent;
        font-weight: bold;
        margin-bottom: 30px;
    }
    .prediction-box {
        padding: 20px;
        border-radius: 15px;
        text-align: center;
        margin: 20px 0;
        animation: fadeIn 0.5s;
    }
    .cat-result {
        background: linear-gradient(135deg, #ffd89b 0%, #19547b 100%);
        color: white;
        font-size: 1.5em;
    }
    .dog-result {
        background: linear-gradient(135deg, #f093fb 0%, #f5576c 100%);
        color: white;
        font-size: 1.5em;
    }
    .confidence-bar {
        height: 30px;
        border-radius: 15px;
        margin: 10px 0;
    }
    .upload-section {
        border: 3px dashed #667eea;
        border-radius: 20px;
        padding: 40px;
        text-align: center;
        background: #f8f9fa;
        transition: all 0.3s;
    }
    .upload-section:hover {
        border-color: #764ba2;
        background: #e8eaf6;
    }
    @keyframes fadeIn {
        from { opacity: 0; transform: translateY(20px); }
        to { opacity: 1; transform: translateY(0); }
    }
    .stats-box {
        background: white;
        padding: 15px;
        border-radius: 10px;
        box-shadow: 0 2px 10px rgba(0,0,0,0.1);
        margin: 10px 0;
    }
</style>
""", unsafe_allow_html=True)

# ═══════════════════════════════════════════════════════════
# LOAD MODEL (Cached)
# ═══════════════════════════════════════════════════════════
@st.cache_resource
def load_classifier_model():
    """Load the trained model"""
    model_path = 'cat_dog_model.h5'
    
    if not os.path.exists(model_path):
        st.error(f"❌ Model file not found: {model_path}")
        st.info("Please train the model first using cat_dog_classifier.py")
        return None
    
    try:
        model = load_model(model_path)
        return model
    except Exception as e:
        st.error(f"❌ Error loading model: {e}")
        return None

# ═══════════════════════════════════════════════════════════
# IMAGE PREPROCESSING
# ═══════════════════════════════════════════════════════════
def preprocess_image(img, target_size=(150, 150)):
    """Preprocess image for model prediction"""
    # Resize image
    img = img.resize(target_size)
    
    # Convert to array
    img_array = image.img_to_array(img)
    
    # Expand dimensions for batch
    img_array = np.expand_dims(img_array, axis=0)
    
    # Normalize
    img_array = img_array / 255.0
    
    return img_array

# ═══════════════════════════════════════════════════════════
# PREDICTION FUNCTION
# ═══════════════════════════════════════════════════════════
def predict_image(model, img):
    """Make prediction on uploaded image"""
    # Preprocess
    processed_img = preprocess_image(img)
    
    # Predict
    prediction = model.predict(processed_img, verbose=0)
    
    # Get confidence
    confidence = float(prediction[0][0])
    
    # Determine class
    if confidence > 0.5:
        predicted_class = "🐶 Dog"
        dog_conf = confidence * 100
        cat_conf = (1 - confidence) * 100
    else:
        predicted_class = "🐱 Cat"
        cat_conf = (1 - confidence) * 100
        dog_conf = confidence * 100
    
    return predicted_class, cat_conf, dog_conf, confidence

# ═══════════════════════════════════════════════════════════
# MAIN APP
# ═══════════════════════════════════════════════════════════
def main():
    # Title
    st.markdown('<h1 class="main-title">🐱 Cat vs Dog Classifier 🐶</h1>', 
                unsafe_allow_html=True)
    
    # Load model
    model = load_classifier_model()
    
    if model is None:
        st.warning("⚠️ Please train the model first or check the model file path.")
        return
    
    st.success("✅ Model loaded successfully!")
    
    # Sidebar
    with st.sidebar:
        st.header("ℹ️ About")
        st.markdown("""
        This app uses a **CNN (Convolutional Neural Network)** 
        trained on cat and dog images to classify your uploaded photos.
        
        **How to use:**
        1. Upload an image of a cat or dog
        2. Click "Predict" button
        3. See the result with confidence score
        
        **Supported formats:** JPG, JPEG, PNG
        
        **Model:** Custom CNN
        **Accuracy:** ~85% (may vary)
        """)
        
        st.markdown("---")
        
        # Model info
        st.subheader("📊 Model Info")
        st.markdown(f"""
        - **Type:** Convolutional Neural Network
        - **Input size:** 150×150 pixels
        - **Layers:** 3 Conv + 3 Pool + 2 Dense
        - **Parameters:** ~3.5M
        """)
        
        st.markdown("---")
        
        # Example images
        st.subheader("🖼️ Sample Predictions")
        col1, col2 = st.columns(2)
        with col1:
            st.markdown("**Cat Example**")
            st.markdown("Confidence threshold: > 0.5 = Dog")
        with col2:
            st.markdown("**Dog Example**")
            st.markdown("Confidence threshold: < 0.5 = Cat")
    
    # Main content
    col1, col2 = st.columns([1, 1])
    
    with col1:
        st.markdown("### 📤 Upload Image")
        
        # File uploader
        uploaded_file = st.file_uploader(
            "Choose an image of a cat or dog...",
            type=['jpg', 'jpeg', 'png'],
            help="Upload a clear image for best results"
        )
    
    with col2:
        st.markdown("### 📸 Or Use Camera")
        camera_image = st.camera_input("Take a photo")
    
    # Process uploaded image
    image_to_process = None
    
    if uploaded_file is not None:
        image_to_process = Image.open(uploaded_file)
        source = "uploaded"
    elif camera_image is not None:
        image_to_process = Image.open(camera_image)
        source = "camera"
    
    if image_to_process is not None:
        st.markdown("---")
        
        # Display image and prediction
        col1, col2 = st.columns([1, 1])
        
        with col1:
            st.markdown("### 🖼️ Your Image")
            st.image(image_to_process, use_column_width=True)
            
            # Image details
            st.markdown(f"""
            <div class="stats-box">
                <b>Image Details:</b><br>
                📐 Size: {image_to_process.size[0]}×{image_to_process.size[1]} px<br>
                🎨 Mode: {image_to_process.mode}<br>
                📁 Source: {source.capitalize()}
            </div>
            """, unsafe_allow_html=True)
        
        with col2:
            st.markdown("### 🤖 Prediction")
            
            # Predict button
            if st.button("🔍 Predict", type="primary", use_container_width=True):
                with st.spinner("🧠 Analyzing image..."):
                    # Add small delay for effect
                    time.sleep(0.5)
                    
                    # Make prediction
                    predicted_class, cat_conf, dog_conf, confidence = predict_image(
                        model, image_to_process
                    )
                
                # Display result
                if "Cat" in predicted_class:
                    result_class = "cat-result"
                    emoji = "🐱"
                else:
                    result_class = "dog-result"
                    emoji = "🐶"
                
                st.markdown(f"""
                <div class="prediction-box {result_class}">
                    <h2>{emoji} {predicted_class}</h2>
                    <h4>Confidence: {max(cat_conf, dog_conf):.1f}%</h4>
                </div>
                """, unsafe_allow_html=True)
                
                # Confidence bars
                st.markdown("#### 📊 Confidence Breakdown")
                
                col_a, col_b = st.columns(2)
                
                with col_a:
                    st.markdown("**Cat Probability**")
                    st.progress(int(cat_conf) / 100)
                    st.markdown(f"{cat_conf:.1f}%")
                
                with col_b:
                    st.markdown("**Dog Probability**")
                    st.progress(int(dog_conf) / 100)
                    st.markdown(f"{dog_conf:.1f}%")
                
                # Raw prediction value
                with st.expander("🔬 Technical Details"):
                    st.markdown(f"""
                    - **Raw prediction value:** {confidence:.6f}
                    - **Threshold:** 0.5
                    - **Prediction logic:** Value > 0.5 = Dog, < 0.5 = Cat
                    - **Model output:** Sigmoid activation
                    """)
    
    # Footer
    st.markdown("---")
    st.markdown(
        "<div style='text-align: center; color: #666;'>"
        "🐾 Cat vs Dog Classifier | Built with TensorFlow & Streamlit | "
        "Upload an image to get started!"
        "</div>",
        unsafe_allow_html=True
    )

# ═══════════════════════════════════════════════════════════
# MULTIPLE IMAGES BATCH PREDICTION (Optional)
# ═══════════════════════════════════════════════════════════
def batch_prediction_section(model):
    """Optional section for batch predictions"""
    st.markdown("---")
    st.markdown("### 📁 Batch Prediction (Multiple Images)")
    
    uploaded_files = st.file_uploader(
        "Upload multiple images",
        type=['jpg', 'jpeg', 'png'],
        accept_multiple_files=True,
        key="batch_upload"
    )
    
    if uploaded_files:
        st.markdown(f"**{len(uploaded_files)} images uploaded**")
        
        if st.button("🔍 Predict All", type="secondary"):
            results = []
            
            progress_bar = st.progress(0)
            
            for i, file in enumerate(uploaded_files):
                img = Image.open(file)
                predicted_class, cat_conf, dog_conf, _ = predict_image(model, img)
                
                results.append({
                    'filename': file.name,
                    'prediction': predicted_class,
                    'cat_confidence': f"{cat_conf:.1f}%",
                    'dog_confidence': f"{dog_conf:.1f}%"
                })
                
                # Update progress
                progress_bar.progress((i + 1) / len(uploaded_files))
            
            # Display results table
            import pandas as pd
            df = pd.DataFrame(results)
            st.dataframe(df, use_container_width=True)
            
            # Summary
            cats = sum(1 for r in results if 'Cat' in r['prediction'])
            dogs = len(results) - cats
            
            col1, col2 = st.columns(2)
            with col1:
                st.metric("🐱 Cats", cats)
            with col2:
                st.metric("🐶 Dogs", dogs)

# ═══════════════════════════════════════════════════════════
# RUN APP
# ═══════════════════════════════════════════════════════════
if __name__ == "__main__":
    main()