#!/usr/bin/env python3
"""
zeroFire - Fire Detection Classification App
AI-powered fire detection using ConvNeXt transfer learning
"""

import streamlit as st
import torch
import torch.nn.functional as F
from PIL import Image
import numpy as np
import plotly.graph_objects as go
import plotly.express as px
from plotly.subplots import make_subplots
import pandas as pd
import sys
import os
import time
from io import BytesIO
import base64

# Add utils to path
sys.path.append('utils')
from model_utils import load_model, FireDetectionClassifier
from data_utils import get_inference_transform, prepare_image_for_inference, check_data_directory

# Page Configuration
st.set_page_config(
    page_title="🔥 zeroFire - Fire Detection System",
    page_icon="🔥",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS for Beautiful UI
st.markdown("""
<style>
    .main-header {
        background: linear-gradient(135deg, #74b9ff 0%, #0984e3 100%);
        padding: 2rem;
        border-radius: 15px;
        text-align: center;
        color: white;
        margin-bottom: 2rem;
        box-shadow: 0 8px 32px rgba(0,0,0,0.1);
    }
    
    .main-header h1 {
        font-size: 3rem;
        margin: 0;
        font-weight: bold;
        text-shadow: 2px 2px 4px rgba(0,0,0,0.3);
    }
    
    .main-header p {
        font-size: 1.2rem;
        margin: 0.5rem 0 0 0;
        opacity: 0.9;
    }
    
    .upload-section {
        background: linear-gradient(135deg, #a8e6cf 0%, #74b9ff 100%);
        color: white;
        padding: 20px 25px;
        border-radius: 15px;
        text-align: center;
        margin-bottom: 20px;
        box-shadow: 0 6px 20px rgba(168, 230, 207, 0.3);
        height: 130px;
        display: flex;
        flex-direction: column;
        justify-content: center;
        line-height: 1.4;
    }
    
    .upload-section h3 {
        font-size: 1.3rem;
        margin: 0 0 8px 0;
        font-weight: bold;
    }
    
    .upload-section p {
        font-size: 0.95rem;
        margin: 0;
        opacity: 0.9;
        line-height: 1.3;
    }
    
    .result-fire {
        background: linear-gradient(135deg, #dc3545 0%, #c82333 100%);
        color: white;
        padding: 15px 20px;
        border-radius: 15px;
        text-align: center;
        margin: 0 0 20px 0;
        box-shadow: 0 6px 20px rgba(220, 53, 69, 0.3);
        height: 100px;
        display: flex;
        flex-direction: column;
        justify-content: center;
    }
    
    .result-fire h2 {
        font-size: 1.2rem;
        margin: 0 0 5px 0;
        line-height: 1.2;
    }
    
    .result-fire p {
        font-size: 0.85rem;
        margin: 0;
        font-weight: bold;
    }
    
    .result-no-fire {
        background: linear-gradient(135deg, #28a745 0%, #20c997 100%);
        color: white;
        padding: 15px 20px;
        border-radius: 15px;
        text-align: center;
        margin: 0 0 20px 0;
        box-shadow: 0 6px 20px rgba(40, 167, 69, 0.3);
        height: 100px;
        display: flex;
        flex-direction: column;
        justify-content: center;
    }
    
    .result-no-fire h2 {
        font-size: 1.2rem;
        margin: 0 0 5px 0;
        line-height: 1.2;
    }
    
    .result-no-fire p {
        font-size: 0.85rem;
        margin: 0;
        font-weight: bold;
    }
    
    .metric-card {
        background: linear-gradient(135deg, #ff9ff3 0%, #f368e0 100%);
        padding: 15px;
        border-radius: 10px;
        text-align: center;
        color: white;
        margin: 10px 0;
        box-shadow: 0 4px 15px rgba(0,0,0,0.2);
    }
    
    .info-card {
        background: linear-gradient(135deg, #ff9ff3 0%, #f368e0 100%);
        color: white;
        padding: 20px;
        border-radius: 15px;
        margin: 15px 0;
        box-shadow: 0 8px 32px rgba(0,0,0,0.1);
    }
    
    .info-card-fire {
        background: linear-gradient(135deg, #ff6b6b 0%, #ee5a52 100%);
        color: white;
        padding: 20px;
        border-radius: 15px;
        margin: 15px 0;
        box-shadow: 0 8px 32px rgba(0,0,0,0.1);
    }
    
    .stButton > button {
        background: linear-gradient(45deg, #74b9ff, #0984e3);
        color: white;
        border: none;
        border-radius: 10px;
        padding: 12px 24px;
        font-weight: bold;
        font-size: 1.1rem;
        transition: all 0.3s ease;
        box-shadow: 0 4px 15px rgba(0,0,0,0.2);
        width: 100%;
    }
    
    .stButton > button:hover {
        transform: translateY(-2px);
        box-shadow: 0 6px 20px rgba(0,0,0,0.3);
    }
    
    .sidebar .stSelectbox > div > div {
        background: linear-gradient(135deg, #74b9ff 0%, #0984e3 100%);
        color: white;
    }
    
    /* Constrain image height to prevent scrolling */
    .stImage > img {
        max-height: 400px;
        width: auto;
        object-fit: contain;
    }
    
    /* File uploader styling - enhanced for mobile */
    .stFileUploader > div {
        border-radius: 15px;
        margin: 10px 0;
    }
    
    .stFileUploader > div > div {
        border: 2px dashed #74b9ff;
        border-radius: 15px;
        padding: 20px;
        background: rgba(116, 185, 255, 0.05);
        transition: all 0.3s ease;
        min-height: 80px;
        display: flex;
        align-items: center;
        justify-content: center;
    }
    
    .stFileUploader > div > div:hover {
        border-color: #0984e3;
        background: rgba(116, 185, 255, 0.1);
        transform: translateY(-2px);
        box-shadow: 0 4px 15px rgba(116, 185, 255, 0.2);
    }
    
    /* Mobile specific styling */
    @media (max-width: 768px) {
        .stFileUploader > div > div {
            min-height: 100px;
            padding: 25px;
            font-size: 1.1rem;
        }
        
        .stFileUploader button {
            padding: 15px 25px;
            font-size: 1.1rem;
            border-radius: 10px;
            background: linear-gradient(135deg, #74b9ff 0%, #0984e3 100%);
            color: white;
            border: none;
            font-weight: bold;
        }
    }
    
    /* Mobile responsive adjustments */
    @media (max-width: 768px) {
        .main-header h1 {
            font-size: 2rem;
        }
        
        .main-header p {
            font-size: 1rem;
        }
        
        .upload-section {
            height: 110px;
            padding: 18px 20px;
        }
        
        .upload-section h3 {
            font-size: 1.1rem;
            margin: 0 0 6px 0;
        }
        
        .upload-section p {
            font-size: 0.9rem;
            line-height: 1.2;
        }
        
        .result-fire h2, .result-no-fire h2 {
            font-size: 1rem;
        }
        
        .result-fire, .result-no-fire {
            height: 80px;
        }
        
        .stImage > img {
            max-height: 300px;
        }
    }
</style>
""", unsafe_allow_html=True)

@st.cache_resource
def load_fire_model():
    """Load the trained fire detection model"""
    model_path = 'models/fire_detection_classifier.pth'
    
    if not os.path.exists(model_path):
        return None, "Model not found. Please train the model first."
    
    try:
        model, model_info = load_model(model_path, device='cpu')
        return model, model_info
    except Exception as e:
        return None, f"Error loading model: {str(e)}"

def get_prediction(image, model, transform):
    """Get prediction from the model"""
    try:
        # Prepare image
        input_tensor = prepare_image_for_inference(image, transform)
        
        # Get prediction
        with torch.no_grad():
            model.eval()
            outputs = model(input_tensor)
            probabilities = F.softmax(outputs, dim=1)
            confidence, predicted = torch.max(probabilities, 1)
            
            # Convert to numpy
            predicted_class = predicted.item()
            confidence_score = confidence.item()
            all_probs = probabilities.squeeze().cpu().numpy()
            
            return predicted_class, confidence_score, all_probs
    
    except Exception as e:
        st.error(f"Error during prediction: {str(e)}")
        return None, None, None

def create_confidence_chart(probabilities, class_names):
    """Create confidence chart using Plotly"""
    fig = go.Figure(data=[
        go.Bar(
            x=class_names,
            y=probabilities,
            marker_color=['#dc3545', '#28a745'],
            text=[f'{p:.1%}' for p in probabilities],
            textposition='auto',
        )
    ])
    
    fig.update_layout(
        title="Fire Detection Confidence",
        xaxis_title="Prediction",
        yaxis_title="Confidence",
        yaxis=dict(range=[0, 1]),
        showlegend=False,
        height=400,
        template="plotly_white"
    )
    
    return fig

def create_safety_metrics_chart(predicted_class, confidence):
    """Create safety metrics visualization"""
    if predicted_class == 0:  # Fire
        danger_level = confidence * 100
        safety_level = (1 - confidence) * 100
        primary_color = '#dc3545'
        status = "FIRE DETECTED"
    else:  # No Fire
        danger_level = (1 - confidence) * 100
        safety_level = confidence * 100
        primary_color = '#28a745'
        status = "NO FIRE"
    
    fig = go.Figure(go.Indicator(
        mode = "gauge+number+delta",
        value = danger_level,
        domain = {'x': [0, 1], 'y': [0, 1]},
        title = {'text': "Fire Risk Level"},
        delta = {'reference': 50},
        gauge = {'axis': {'range': [None, 100]},
                'bar': {'color': primary_color},
                'steps' : [
                    {'range': [0, 25], 'color': "lightgray"},
                    {'range': [25, 50], 'color': "yellow"},
                    {'range': [50, 75], 'color': "orange"},
                    {'range': [75, 100], 'color': "red"}],
                'threshold' : {'line': {'color': "red", 'width': 4},
                              'thickness': 0.75, 'value': 90}}))
    
    fig.update_layout(height=400)
    return fig

def analyze_fire_risk(predicted_class, confidence):
    """Analyze fire risk and provide recommendations"""
    if predicted_class == 0:  # Fire detected
        risk_level = confidence * 100
        
        if risk_level >= 90:
            return {
                'level': 'CRITICAL',
                'color': '#dc3545',
                'icon': '🚨',
                'message': 'IMMEDIATE ACTION REQUIRED',
                'recommendations': [
                    'Activate fire suppression system immediately',
                    'Evacuate the area',
                    'Call emergency services UAE 997',
                    'Shut down affected equipment if safe to do so',
                    'Monitor surrounding areas for spread'
                ]
            }
        elif risk_level >= 75:
            return {
                'level': 'HIGH',
                'color': '#fd7e14',
                'icon': '⚠️',
                'message': 'HIGH FIRE RISK DETECTED',
                'recommendations': [
                    'Investigate the area immediately',
                    'Prepare fire suppression systems',
                    'Alert security personnel',
                    'Consider equipment shutdown',
                    'Increase monitoring frequency'
                ]
            }
        else:
            return {
                'level': 'MODERATE',
                'color': '#ffc107',
                'icon': '🔶',
                'message': 'POSSIBLE FIRE DETECTED',
                'recommendations': [
                    'Verify with additional sensors',
                    'Send personnel to investigate',
                    'Check equipment temperatures',
                    'Review recent maintenance logs',
                    'Maintain heightened awareness'
                ]
            }
    else:  # No fire
        safety_level = confidence * 100
        
        if safety_level >= 95:
            return {
                'level': 'SAFE',
                'color': '#28a745',
                'icon': '✅',
                'message': 'NORMAL OPERATION',
                'recommendations': [
                    'Continue normal operations',
                    'Maintain regular monitoring',
                    'Keep fire suppression systems ready',
                    'Perform scheduled maintenance',
                    'Review safety protocols periodically'
                ]
            }
        else:
            return {
                'level': 'CAUTION',
                'color': '#17a2b8',
                'icon': '⚠️',
                'message': 'MONITOR CLOSELY',
                'recommendations': [
                    'Increase monitoring frequency',
                    'Check for unusual conditions',
                    'Verify sensor functionality',
                    'Review environmental factors',
                    'Maintain readiness for action'
                ]
            }

# display_fire_safety_checklist function removed - content moved to right column

def main():
    """Main application function"""
    # Header
    st.markdown("""
    <div class="main-header">
        <h1>🔥 zeroFire</h1>
        <p>AI-Powered Fire Detection System for Data Centers</p>
    </div>
    """, unsafe_allow_html=True)
    
    # Sidebar
    with st.sidebar:
        st.markdown("### 🔥 Fire Detection System")
        st.markdown("---")
        
        # Model status
        model, model_info = load_fire_model()
        
        if model is None:
            st.error("❌ Model not available")
            st.info("Train the model first using: `python train_fire_detection.py`")
            return
        else:
            st.success("✅ Model loaded successfully")
            if isinstance(model_info, dict):
                accuracy = model_info.get('best_acc', 'Unknown')
                if accuracy != 'Unknown':
                    # Format accuracy to 2 decimal places
                    accuracy_formatted = f"{float(accuracy):.2f}%"
                    st.info(f"📊 Model: ConvNeXt Large")
                    st.info(f"🎯 Accuracy: {accuracy_formatted}")
                    st.info(f"🔄 Transfer Learning: Food→Fire")
                    st.info(f"⚡ Precision: High-recall optimized")
                else:
                    st.info("📊 Model Accuracy: Unknown")
        
        st.markdown("---")
        
        # Settings
        st.markdown("### ⚙️ Settings")
        confidence_threshold = st.slider(
            "Confidence Threshold",
            min_value=0.0,
            max_value=1.0,
            value=0.5,
            step=0.05,
            help="Minimum confidence required for fire detection"
        )
        
        show_details = st.checkbox("Show detailed analysis", value=True)
        
        st.markdown("---")
        st.markdown("### 📊 Data Center Status")
        check_data_directory('data')
    
    # Main content
    col1, col2 = st.columns([2, 1])
    
    with col1:
        # Upload section
        st.markdown("""
        <div class="upload-section">
            <h3>📸 Fire Detection Input</h3>
            <p>Upload an image or take a photo to detect fire or smoke<br>in your data center</p>
        </div>
        """, unsafe_allow_html=True)
        
        # Mobile-friendly file uploader
        st.markdown("**📱 Take Photo or Upload Image**")
        
        # Clear mobile instructions
        st.markdown("""
        <div style="background: #e3f2fd; padding: 15px; border-radius: 10px; margin: 15px 0; border-left: 4px solid #2196f3;">
            <p style="margin: 0; color: #1976d2; font-weight: bold;">📱 Mobile Users:</p>
            <p style="margin: 5px 0 0 0; color: #1976d2;">
                Tap "Browse files" below → Select <strong>"Camera"</strong> to take a photo<br>
                Or select <strong>"Photos"</strong> to choose from gallery
            </p>
        </div>
        """, unsafe_allow_html=True)
        
        # Troubleshooting info
        with st.expander("🔧 Camera not working? Click here for help"):
            st.markdown("""
            **If you don't see the Camera option:**
            
            1. **✅ Check your browser:**
               - Use Chrome, Safari, or Firefox (latest versions)
               - Edge or other browsers may not support camera
            
            2. **🔐 Ensure secure connection:**
               - Camera requires HTTPS (✅ Hugging Face uses HTTPS)
               - Local development requires HTTPS for camera access
            
            3. **📱 Mobile device requirements:**
               - iOS: Safari 11+ or Chrome 64+
               - Android: Chrome 53+ or Firefox 68+
               - Some older devices may not support camera
            
            4. **🛠️ Try these steps:**
               - Refresh the page and try again
               - Clear browser cache and cookies
               - Try a different browser
               - Check if camera works on other websites
            
            5. **🔄 Alternative options:**
               - Take photo with your camera app first
               - Then select "Photos" to upload the saved image
               - Or use the desktop version for file upload
            """)
            
            st.info("💡 **Note**: Camera access depends on your browser and device. If it doesn't work, you can still upload photos from your gallery!")
        
        # Standard Streamlit file uploader
        uploaded_file = st.file_uploader(
            "Browse files",
            type=['jpg', 'jpeg', 'png'],
            help="📱 Mobile: Tap to see Camera and Photos options | 💻 Desktop: Click to browse files"
        )
        
        # Process the uploaded image
        image = None
        image_source = "uploaded"
        
        if uploaded_file is not None:
            image = Image.open(uploaded_file)
            image_source = "uploaded"
        
        # Process the image (whether uploaded or from camera)
        if image is not None:
            # Display the image with appropriate caption
            st.image(image, caption="📸 Your Image", use_column_width=True)
            
            # Get prediction
            transform = get_inference_transform()
            predicted_class, confidence_score, all_probs = get_prediction(image, model, transform)
            
            if predicted_class is not None:
                class_names = ['Fire', 'No Fire']
                predicted_label = class_names[predicted_class]
                
                # Fire detection result moved to right column
                
                # Technical details (keep only this in left column)
                with st.expander("🔬 Technical Details"):
                    st.markdown(f"""
                    **Prediction Details:**
                    - Predicted Class: {predicted_label}
                    - Confidence Score: {confidence_score:.4f}
                    - Fire Probability: {all_probs[0]:.4f}
                    - No-Fire Probability: {all_probs[1]:.4f}
                    - Threshold: {confidence_threshold:.2f}
                    """)
    
    with col2:
        # Display fire detection result first
        if 'predicted_class' in locals():
            # Display fire detection result box
            if predicted_class == 0:  # Fire
                st.markdown(f"""
                <div class="result-fire">
                    <h2>🚨 FIRE DETECTED - Confidence: {confidence_score:.1%}</h2>
                    <p>IMMEDIATE ACTION REQUIRED</p>
                </div>
                """, unsafe_allow_html=True)
            else:  # No Fire
                st.markdown(f"""
                <div class="result-no-fire">
                    <h2>✅ NO FIRE DETECTED - Confidence: {confidence_score:.1%}</h2>
                    <p>Normal Operation</p>
                </div>
                """, unsafe_allow_html=True)
        
        # Quick metrics
        st.markdown("### 📊 Quick Metrics")
        
        if 'predicted_class' in locals():
            # Create three columns for metrics in a row
            metric_col1, metric_col2, metric_col3 = st.columns(3)
            
            with metric_col1:
                # Fire risk gauge
                risk_percentage = all_probs[0] * 100
                st.metric(
                    label="Fire Risk",
                    value=f"{risk_percentage:.1f}%"
                )
            
            with metric_col2:
                # Safety score
                safety_score = all_probs[1] * 100
                st.metric(
                    label="Safety Score",
                    value=f"{safety_score:.1f}%"
                )
            
            with metric_col3:
                # Model status instead of redundant confidence
                status = "FIRE ALERT" if predicted_class == 0 else "NORMAL"
                st.metric(
                    label="Status",
                    value=status
                )
            
                         # Charts removed to eliminate confusion and redundancy
            
            # Detailed analysis moved from left column
            if show_details:
                st.markdown("### 📋 Detailed Analysis")
                
                analysis = analyze_fire_risk(predicted_class, confidence_score)
                
                # Use red styling only for fire detection
                card_class = "info-card-fire" if predicted_class == 0 else "info-card"
                
                st.markdown(f"""
                <div class="{card_class}">
                    <h3>{analysis['icon']} Risk Level: {analysis['level']}</h3>
                    <p><strong>{analysis['message']}</strong></p>
                </div>
                """, unsafe_allow_html=True)
                
                st.markdown("#### 🎯 Recommended Actions:")
                for rec in analysis['recommendations']:
                    st.markdown(f"- {rec}")
                
                # Fire Safety Checklist moved here
                st.markdown("---")
                st.markdown("""
                <div class="info-card">
                    <h3>🔥 Fire Safety Checklist</h3>
                    <p>Essential fire safety measures for data centers:</p>
                </div>
                """, unsafe_allow_html=True)
                
                safety_items = [
                    "🔥 **Fire Detection Systems** - Smoke, heat, and flame detectors",
                    "💨 **Suppression Systems** - Clean agent, water mist, or CO2",
                    "🌡️ **Temperature Monitoring** - Continuous thermal monitoring",
                    "⚡ **Electrical Safety** - Arc fault and ground fault protection",
                    "🚪 **Emergency Exits** - Clear and well-marked escape routes",
                    "📋 **Emergency Procedures** - Staff training and evacuation plans",
                    "🔧 **Equipment Maintenance** - Regular inspection and testing",
                    "📞 **Emergency Contacts** - Quick access to fire department",
                    "🎯 **Response Plans** - Pre-defined actions for different scenarios",
                    "📊 **Documentation** - Incident logging and safety records"
                ]
                
                for item in safety_items:
                    st.markdown(f"- {item}")
        
        # Additional info moved to bottom
    
    # Footer
    st.markdown("---")
    
    # Emergency Contacts only (Fire Safety Checklist moved to right column)
    st.markdown("""
    <div class="info-card" style="padding: 15px;">
        <h3 style="margin: 0 0 10px 0; text-align: center;">🚨 Emergency Contacts</h3>
        <div style="display: flex; justify-content: space-around; flex-wrap: wrap; gap: 15px;">
            <span><strong>Fire Dept:</strong> UAE 997</span>
            <span><strong>Security:</strong> [Your Number]</span>
            <span><strong>Facilities:</strong> [Your Number]</span>
            <span><strong>IT Ops:</strong> [Your Number]</span>
        </div>
    </div>
    """, unsafe_allow_html=True)
    
    st.markdown("---")
    st.markdown("""
    <div style="text-align: center; padding: 20px; background: linear-gradient(135deg, #a8e6cf 0%, #74b9ff 100%); border-radius: 15px; color: white;">
        <p>🔥 zeroFire - AI-Powered Fire Detection System</p>
        <p>Protecting your data center with advanced machine learning</p>
    </div>
    """, unsafe_allow_html=True)

if __name__ == "__main__":
    main()