#!/usr/bin/env python3
"""
ACCEPTIN - Telecom Site Quality Classification App
AI-powered telecom site inspection 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, TelecomClassifier
from data_utils import get_inference_transform, prepare_image_for_inference, check_data_directory

# Page Configuration
st.set_page_config(
    page_title="📡 ACCEPTIN - Telecom Site Inspector",
    page_icon="📡",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS for Beautiful UI
st.markdown("""
<style>
    .main-header {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 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, #56ab2f 0%, #a8e6cf 100%);
        color: white;
        padding: 20px;
        border-radius: 15px;
        text-align: center;
        margin-bottom: 20px;
        box-shadow: 0 6px 20px rgba(86, 171, 47, 0.3);
    }
    
    .result-good {
        background: linear-gradient(135deg, #28a745 0%, #20c997 100%);
        color: white;
        padding: 20px;
        border-radius: 15px;
        text-align: center;
        margin: 20px 0;
        box-shadow: 0 6px 20px rgba(40, 167, 69, 0.3);
    }
    
    .result-bad {
        background: linear-gradient(135deg, #dc3545 0%, #e83e8c 100%);
        color: white;
        padding: 20px;
        border-radius: 15px;
        text-align: center;
        margin: 20px 0;
        box-shadow: 0 6px 20px rgba(220, 53, 69, 0.3);
    }
    
    .metric-card {
        background: linear-gradient(135deg, #4facfe 0%, #00f2fe 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, #667eea 0%, #764ba2 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, #667eea, #764ba2);
        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, #667eea 0%, #764ba2 100%);
        color: white;
    }
</style>
""", unsafe_allow_html=True)

@st.cache_resource
def load_telecom_model():
    """Load the trained telecom classification model"""
    model_path = 'models/telecom_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="Classification Confidence",
        xaxis_title="Site Quality",
        yaxis_title="Confidence",
        yaxis=dict(range=[0, 1]),
        showlegend=False,
        height=400,
        template="plotly_white"
    )
    
    return fig

def create_quality_metrics_chart(predicted_class, confidence):
    """Create quality metrics visualization"""
    if predicted_class == 1:  # Good
        quality_score = confidence * 100
        color = '#28a745'
        status = 'ACCEPTED'
    else:  # Bad
        quality_score = (1 - confidence) * 100
        color = '#dc3545'
        status = 'REJECTED'
    
    fig = go.Figure(go.Indicator(
        mode="gauge+number+delta",
        value=quality_score,
        domain={'x': [0, 1], 'y': [0, 1]},
        title={'text': f"Quality Score<br><span style='font-size:0.8em;color:{color}'>{status}</span>"},
        delta={'reference': 80},
        gauge={
            'axis': {'range': [None, 100]},
            'bar': {'color': color},
            'steps': [
                {'range': [0, 50], 'color': "lightgray"},
                {'range': [50, 80], 'color': "yellow"},
                {'range': [80, 100], 'color': "lightgreen"}
            ],
            'threshold': {
                'line': {'color': "red", 'width': 4},
                'thickness': 0.75,
                'value': 90
            }
        }
    ))
    
    fig.update_layout(height=400)
    return fig

def analyze_site_quality(predicted_class, confidence):
    """Analyze site quality and provide detailed feedback"""
    class_names = ['Bad', 'Good']
    predicted_label = class_names[predicted_class]
    
    if predicted_class == 1:  # Good site
        analysis = {
            'status': 'ACCEPTED ✅',
            'color': 'result-good',
            'icon': '✅',
            'message': 'Site installation meets quality standards',
            'details': [
                '✅ Cable assembly appears properly organized',
                '✅ Equipment installation looks correct',
                '✅ Overall site organization is acceptable',
                '✅ No obvious safety violations detected'
            ],
            'recommendations': [
                '🔍 Verify all labels are clearly readable',
                '🔧 Double-check all card installations',
                '📋 Complete final inspection checklist',
                '📸 Document final installation state'
            ]
        }
    else:  # Bad site
        analysis = {
            'status': 'REJECTED ❌',
            'color': 'result-bad',
            'icon': '❌',
            'message': 'Site installation requires attention',
            'details': [
                '❌ Cable organization may need improvement',
                '❌ Equipment installation issues detected',
                '❌ Site organization below standards',
                '❌ Potential safety or quality concerns'
            ],
            'recommendations': [
                '🔧 Reorganize cable routing and bundling',
                '🔍 Check all card installations and seating',
                '🏷️ Verify all labels are present and readable',
                '⚠️ Address any safety violations',
                '📋 Complete corrective actions before acceptance'
            ]
        }
    
    analysis['confidence'] = confidence
    analysis['predicted_label'] = predicted_label
    
    return analysis

def display_inspection_checklist():
    """Display telecom site inspection checklist"""
    st.markdown("""
    <div class="info-card">
        <h3>📋 Telecom Site Inspection Checklist</h3>
        <p>Use this checklist to ensure comprehensive site evaluation:</p>
    </div>
    """, unsafe_allow_html=True)
    
    checklist_items = {
        "Cable Assembly": [
            "Cables properly routed and bundled",
            "No loose or hanging cables",
            "Proper cable management systems used",
            "Cable routing follows standards"
        ],
        "Card Installation": [
            "All required cards present",
            "Cards properly seated and secured",
            "No missing or damaged cards",
            "Card configurations correct"
        ],
        "Labeling": [
            "All equipment properly labeled",
            "Labels clearly readable",
            "Label placement follows standards",
            "No missing identification tags"
        ],
        "Safety & Organization": [
            "Safety covers properly installed",
            "Grounding connections secure",
            "Warning signs present where required",
            "Overall rack organization acceptable"
        ]
    }
    
    for category, items in checklist_items.items():
        st.subheader(f"🔍 {category}")
        for item in items:
            st.write(f"• {item}")

def main():
    """Main application function"""
    # Header
    st.markdown("""
    <div class="main-header">
        <h1>📡 ACCEPTIN</h1>
        <p>AI-Powered Telecom Site Quality Inspector</p>
    </div>
    """, unsafe_allow_html=True)
    
    # Sidebar
    st.sidebar.title("🛠️ Controls")
    
    # Load model
    model, model_info = load_telecom_model()
    
    if model is None:
        st.error(f"❌ {model_info}")
        st.info("Please train the model first using: `python train_telecom.py`")
        return
    
    # Model info in sidebar
    st.sidebar.success("✅ Model loaded successfully")
    if isinstance(model_info, dict):
        st.sidebar.write(f"**Accuracy:** {model_info.get('best_acc', 'Unknown')}")
        st.sidebar.write(f"**Architecture:** ConvNeXt Large")
        st.sidebar.write(f"**Task:** Binary Classification")
    
    # Main content
    col1, col2 = st.columns([1, 1])
    
    with col1:
        st.markdown("""
        <div class="upload-section">
            <h3>📤 Upload or Capture Telecom Site Image</h3>
            <p>Upload an image or take a photo of the telecom site for quality inspection</p>
        </div>
        """, unsafe_allow_html=True)
        
        # Input method selection
        input_method = st.selectbox(
            "Choose how to provide the telecom site image:",
            ["📁 Upload from device", "📷 Take photo with camera"],
            help="Select whether to upload an existing image or take a new photo"
        )
        
        image = None
        if input_method == "📁 Upload from device":
            uploaded_file = st.file_uploader(
                "Choose an image...",
                type=['jpg', 'jpeg', 'png', 'bmp', 'tiff'],
                help="Upload a clear image of the telecom site installation"
            )
            if uploaded_file is not None:
                image = Image.open(uploaded_file)
        elif input_method == "📷 Take photo with camera":
            camera_photo = st.camera_input("Take a photo of the telecom site")
            if camera_photo is not None:
                image = Image.open(camera_photo)
        
        if image is not None:
            # Display uploaded or captured image
            st.image(image, caption="Telecom Site Image", use_column_width=True)
            with st.spinner("Analyzing site quality..."):
                # Get prediction
                transform = get_inference_transform()
                predicted_class, confidence, probabilities = get_prediction(
                    image, model, transform
                )
                if predicted_class is not None:
                    # Confidence thresholding for OOD detection
                    if max(probabilities) <= 0.8:
                        st.warning("⚠️ This image does not appear to be a telecom site. Please upload a valid telecom site photo.")
                        st.session_state.prediction_results = None
                    else:
                        # Store results in session state
                        st.session_state.prediction_results = {
                            'predicted_class': predicted_class,
                            'confidence': confidence,
                            'probabilities': probabilities,
                            'analysis': analyze_site_quality(predicted_class, confidence)
                        }
                        st.success("✅ Analysis complete!")
    
    with col2:
        if hasattr(st.session_state, 'prediction_results'):
            results = st.session_state.prediction_results
            analysis = results['analysis']
            # Display main result
            st.markdown(f"""
            <div class="{analysis['color']}">
                <h2>{analysis['icon']} {analysis['status']}</h2>
                <h3>{analysis['message']}</h3>
                <p><strong>Confidence:</strong> {analysis['confidence']:.1%}</p>
            </div>
            """, unsafe_allow_html=True)
            # Confidence chart
            st.plotly_chart(
                create_confidence_chart(
                    results['probabilities'],
                    ['Bad', 'Good']
                ),
                use_container_width=True
            )
            # Quality metrics
            st.plotly_chart(
                create_quality_metrics_chart(
                    results['predicted_class'],
                    results['confidence']
                ),
                use_container_width=True
            )
    
    # Detailed analysis section
    if hasattr(st.session_state, 'prediction_results'):
        st.markdown("---")
        st.header("📊 Detailed Analysis")
        
        analysis = st.session_state.prediction_results['analysis']
        
        col1, col2 = st.columns([1, 1])
        
        with col1:
            st.subheader("🔍 Quality Assessment")
            for detail in analysis['details']:
                st.write(detail)
        
        with col2:
            st.subheader("💡 Recommendations")
            for recommendation in analysis['recommendations']:
                st.write(recommendation)
    
    # Tabs for additional features
    st.markdown("---")
    tab1, tab2, tab3 = st.tabs(["📋 Inspection Checklist", "📈 Training Data", "ℹ️ About"])
    
    with tab1:
        display_inspection_checklist()
    
    with tab2:
        st.header("📈 Training Data Overview")
        
        # Check data directory
        data_counts = check_data_directory('data')
        
        if data_counts:
            # Create DataFrame for visualization
            data_list = []
            for split, counts in data_counts.items():
                for class_name, count in counts.items():
                    data_list.append({
                        'Split': split.title(),
                        'Class': class_name.title(),
                        'Count': count
                    })
            
            df = pd.DataFrame(data_list)
            
            # Create bar chart
            fig = px.bar(
                df,
                x='Class',
                y='Count',
                color='Split',
                title='Training Data Distribution',
                barmode='group'
            )
            st.plotly_chart(fig, use_container_width=True)
            
            # Display summary table
            st.subheader("📊 Data Summary")
            st.dataframe(df.pivot(index='Class', columns='Split', values='Count'))
        else:
            st.info("No training data found. Please prepare your dataset in the `data/` directory.")
    
    with tab3:
        st.header("ℹ️ About ACCEPTIN")
        st.markdown("""
        **ACCEPTIN** is an AI-powered telecom site quality inspection system that uses computer vision 
        to automatically classify telecom installations as "good" or "bad" based on visual criteria.
        
        ### 🎯 Key Features:
        - **Transfer Learning**: Leverages pre-trained ConvNeXt model (197M parameters)
        - **Binary Classification**: Classifies sites as good/bad with confidence scores
        - **Quality Assessment**: Evaluates cable assembly, card installation, and labeling
        - **Real-time Analysis**: Instant feedback on site quality
        
        ### 🔧 Technical Details:
        - **Model**: ConvNeXt Large with custom classification head
        - **Training**: Transfer learning from food detection model
        - **Input**: 224x224 RGB images
        - **Output**: Binary classification with confidence scores
        
        ### 📊 Quality Criteria:
        - Cable assembly and routing
        - Card installation and labeling
        - General organization and safety
        - Compliance with telecom standards
        
        ### 🚀 Usage:
        1. Upload telecom site image
        2. Click "Analyze Site Quality"
        3. Review results and recommendations
        4. Use inspection checklist for verification
        """)

if __name__ == "__main__":
    main()