import streamlit as st
import cv2
from PIL import Image, ImageEnhance
import numpy as np
import time
from skimage.metrics import structural_similarity as ssim
import base64
from datetime import datetime
import torch

# Load pre-trained YOLOv5 model for object detection
@st.cache_resource
def load_yolo_model():
    model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
    return model

def load_css():
    st.markdown("""
        <style>
        @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600&display=swap');
        
        .stApp {
            background: linear-gradient(135deg, #1a1a1a 0%, #2d2d2d 100%);
            font-family: 'Inter', sans-serif;
            color: #e0e0e0;
        }
        
        .main {
            padding: 2rem;
            max-width: 1200px;
            margin: 0 auto;
        }
        
        .stButton>button {
            background: linear-gradient(135deg, #2196F3 0%, #1976D2 100%);
            color: white;
            padding: 0.75rem 1.5rem;
            border-radius: 10px;
            border: none;
            box-shadow: 0 4px 6px rgba(0,0,0,0.2);
            transition: all 0.3s ease;
            font-weight: 500;
            letter-spacing: 0.5px;
        }
        
        .stButton>button:hover {
            transform: translateY(-2px);
            box-shadow: 0 6px 12px rgba(0,0,0,0.3);
        }
        
        .upload-container {
            background: #2d2d2d;
            border-radius: 15px;
            padding: 1.5rem;
            box-shadow: 0 4px 6px rgba(0,0,0,0.2);
            transition: all 0.3s ease;
            margin-bottom: 1rem;
        }
        
        .upload-container:hover {
            box-shadow: 0 6px 12px rgba(0,0,0,0.3);
        }
        
        .upload-box {
            border: 2px dashed #404040;
            border-radius: 12px;
            padding: 2rem;
            text-align: center;
            background: #333333;
            transition: all 0.3s ease;
            cursor: pointer;
        }
        
        .upload-box:hover {
            border-color: #2196F3;
            background: #383838;
        }
        
        .results-container {
            background: #2d2d2d;
            border-radius: 15px;
            padding: 2rem;
            box-shadow: 0 4px 6px rgba(0,0,0,0.2);
            color: #e0e0e0;
        }
        
        .metric-card {
            background: #333333;
            border-radius: 10px;
            padding: 1rem;
            margin: 0.5rem 0;
            border-left: 4px solid #2196F3;
            color: #e0e0e0;
        }
        
        .stProgress > div > div {
            background: linear-gradient(90deg, #2196F3, #64B5F6);
            border-radius: 10px;
        }
        
        @keyframes pulse {
            0% { opacity: 1; }
            50% { opacity: 0.5; }
            100% { opacity: 1; }
        }
        
        .loading {
            animation: pulse 1.5s infinite;
        }
        </style>
    """, unsafe_allow_html=True)

def enhance_image(image):
    """
    Basic image enhancement with default settings
    """
    enhancer = ImageEnhance.Brightness(image)
    image = enhancer.enhance(1.0)
    enhancer = ImageEnhance.Contrast(image)
    image = enhancer.enhance(1.0)
    enhancer = ImageEnhance.Sharpness(image)
    image = enhancer.enhance(1.0)
    return image

def compare_images(img1, img2, progress_bar):
    """
    Compare two images and return the processed image, similarity score, and difference percentage
    """
    try:
        progress_bar.progress(0)
        
        # Convert images to numpy arrays and ensure same size
        img1 = np.array(img1.resize(img2.size))
        img2 = np.array(img2)
        progress_bar.progress(20)
        
        # Normalize images
        img1 = cv2.normalize(img1, None, 0, 255, cv2.NORM_MINMAX)
        img2 = cv2.normalize(img2, None, 0, 255, cv2.NORM_MINMAX)
        
        # Convert to grayscale
        gray1 = cv2.cvtColor(img1, cv2.COLOR_RGB2GRAY)
        gray2 = cv2.cvtColor(img2, cv2.COLOR_RGB2GRAY)
        progress_bar.progress(40)
        
        # Calculate SSIM
        score, diff = ssim(gray1, gray2, full=True)
        progress_bar.progress(60)
        
        # Generate heatmap
        diff = (diff * 255).astype(np.uint8)
        heatmap = cv2.applyColorMap(diff, cv2.COLORMAP_JET)
        progress_bar.progress(80)
        
        # Highlight differences in red color
        diff_mask = cv2.absdiff(gray1, gray2)
        diff_mask = cv2.cvtColor(diff_mask, cv2.COLOR_GRAY2RGB)
        diff_mask[np.where((diff_mask == [255, 255, 255]).all(axis=2))] = [0, 0, 255]  # Red color for differences
        
        # Combine original image with difference mask
        result_img = cv2.addWeighted(img1, 0.7, diff_mask, 0.3, 0)
        
        # Calculate pixel-wise differences
        diff_percentage = (np.count_nonzero(diff_mask[:, :, 2] > 0) / (diff_mask.shape[0] * diff_mask.shape[1])) * 100
        
        # Ensure that the difference percentage is consistent with the similarity score
        diff_percentage = 100 - (score * 100)
        
        progress_bar.progress(100)
        
        return result_img, score, diff_percentage, heatmap
        
    except Exception as e:
        st.error(f"Error comparing images: {str(e)}")
        return None, 0, 0, None

def detect_objects(image, model):
    """
    Perform object detection on the image using YOLOv5
    """
    try:
        results = model(image)
        results_df = results.pandas().xyxy[0]
        return results_df
    except Exception as e:
        st.error(f"Error in object detection: {str(e)}")
        return None

def draw_object_boxes(image, objects_df):
    """
    Draw bounding boxes on the image for detected objects
    """
    for _, row in objects_df.iterrows():
        xmin, ymin, xmax, ymax, confidence, class_name = int(row['xmin']), int(row['ymin']), int(row['xmax']), int(row['ymax']), row['confidence'], row['name']
        # Draw bounding box
        cv2.rectangle(image, (xmin, ymin), (xmax, ymax), (0, 255, 0), 2)
        # Add label
        cv2.putText(image, f"{class_name} {confidence:.2f}", (xmin, ymin - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
    return image

def main():
    load_css()
    
    # Initialize session state for results
    if "results" not in st.session_state:
        st.session_state.results = None
    
    # Load YOLOv5 model
    yolo_model = load_yolo_model()
    
    # App header
    st.markdown("""
        <div style='text-align: center; margin-bottom: 2rem; background: linear-gradient(135deg, #2196F3 0%, #1976D2 100%); padding: 2rem; border-radius: 15px; color: white;'>
            <h1 style='margin: 0;'>🔍 Image Comparison Tool</h1>
            <p style='margin: 1rem 0 0 0; opacity: 0.9;'>Compare images, highlight differences, and detect objects</p>
        </div>
    """, unsafe_allow_html=True)
    
    # Main content for image upload and display
    st.markdown("<div class='upload-container'>", unsafe_allow_html=True)
    st.markdown("### 📁 Upload Images")
    
    col1, col2 = st.columns(2)
    
    # Reference Image Upload
    with col1:
        reference_image = st.file_uploader(
            "Drop or select reference image",
            type=["jpg", "jpeg", "png"],
            key="reference"
        )
        if reference_image:
            img1 = Image.open(reference_image)
            img1 = enhance_image(img1)
            st.image(img1, caption="Reference Image", use_column_width=True)
            # Clear previous results when a new image is uploaded
            st.session_state.results = None
    
    # New Image Upload
    with col2:
        new_image = st.file_uploader(
            "Drop or select comparison image",
            type=["jpg", "jpeg", "png"],
            key="new"
        )
        if new_image:
            img2 = Image.open(new_image)
            img2 = enhance_image(img2)
            st.image(img2, caption="Comparison Image", use_column_width=True)
            # Clear previous results when a new image is uploaded
            st.session_state.results = None
    
    st.markdown("</div>", unsafe_allow_html=True)
    
    # Sidebar for results and download
    st.sidebar.markdown("### 🎯 Analysis Results")
    
    if reference_image and new_image:
        compare_button = st.sidebar.button("🔍 Analyze Images", use_container_width=True)
        
        if compare_button or st.session_state.results:
            if not st.session_state.results:
                with st.spinner("Processing images..."):
                    progress_bar = st.sidebar.progress(0)
                    
                    start_time = time.time()
                    result_img, score, diff_percentage, heatmap = compare_images(img1, img2, progress_bar)
                    processing_time = time.time() - start_time
                    
                    # Perform object detection
                    objects_df = detect_objects(result_img, yolo_model)
                    
                    # Draw bounding boxes on the analyzed image
                    if objects_df is not None:
                        result_img = draw_object_boxes(result_img, objects_df)
                    
                    # Store results in session state
                    st.session_state.results = {
                        "result_img": result_img,
                        "heatmap": heatmap,
                        "score": score,
                        "diff_percentage": diff_percentage,
                        "processing_time": processing_time,
                        "objects_df": objects_df
                    }
            
            # Display analyzed image (processed image with differences highlighted) in sidebar
            st.sidebar.image(st.session_state.results["result_img"], caption="Analyzed Image (Differences Highlighted)", use_column_width=True)
            
            # Display heatmap in sidebar
            st.sidebar.image(st.session_state.results["heatmap"], caption="Heatmap", use_column_width=True)
            
            # Display metrics in sidebar
            st.sidebar.markdown("### 📊 Metrics")
            st.sidebar.markdown(f"""
                <div class='metric-card'>
                    <h4>Similarity Score</h4>
                    <h2 style='color: #2196F3'>{st.session_state.results["score"]:.2%}</h2>
                </div>
            """, unsafe_allow_html=True)
            
            st.sidebar.markdown(f"""
                <div class='metric-card'>
                    <h4>Difference Detected</h4>
                    <h2 style='color: #2196F3'>{st.session_state.results["diff_percentage"]:.2f}%</h2>
                </div>
            """, unsafe_allow_html=True)
            
            st.sidebar.markdown(f"""
                <div class='metric-card'>
                    <h4>Processing Time</h4>
                    <h2 style='color: #2196F3'>{st.session_state.results["processing_time"]:.2f}s</h2>
                </div>
            """, unsafe_allow_html=True)
            
            # Display detected objects
            if st.session_state.results["objects_df"] is not None:
                st.sidebar.markdown("### 🔍 Detected Objects")
                st.sidebar.dataframe(st.session_state.results["objects_df"])
            
            # Download analyzed image
            st.sidebar.markdown("### 📥 Download Analyzed Image")
            st.sidebar.download_button(
                "Download Analyzed Image",
                data=cv2.imencode('.png', cv2.cvtColor(st.session_state.results["result_img"], cv2.COLOR_RGB2BGR))[1].tobytes(),
                file_name=f"analyzed_image_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png",
                mime="image/png"
            )
    
    # Footer
    st.markdown("""
        <div style='text-align: center; margin-top: 2rem; padding: 1rem; background: #2d2d2d; border-radius: 10px; box-shadow: 0 4px 6px rgba(0,0,0,0.2);'>
            <p style='color: #888; margin: 0;'>Built with ❤️ using Streamlit | Last updated: December 2024</p>
            <p style='color: #888; font-size: 0.9em; margin: 0.5rem 0 0 0;'>Image Comparison Tool v1.0</p>
        </div>
    """, unsafe_allow_html=True)

if __name__ == "__main__":
    main()