import streamlit as st
import cv2
import numpy as np
from PIL import Image
from rembg import remove
import matplotlib.pyplot as plt
import io
import os
import tempfile

# Set page configuration
st.set_page_config(
    page_title="SpotRadar",
    page_icon="🍎",
    layout="wide"
)

# Page title and description
st.title("🍎 SpotRadar - Phase 1 Demo")
st.write("Upload an image of a fruit to analyze dark spots and blemishes.")

# Function to remove the background using rembg
def remove_background(input_image):
    # Convert PIL Image to bytes
    img_byte_arr = io.BytesIO()
    input_image.save(img_byte_arr, format='PNG')
    input_bytes = img_byte_arr.getvalue()
    
    # Remove background
    output_bytes = remove(input_bytes)
    
    # Return as PIL Image
    return Image.open(io.BytesIO(output_bytes))

# Function to create a white background version of segmented spots
def create_white_bg_spots(mask, original_image):
    # Create a white background (all 255s)
    white_bg = np.ones_like(original_image) * 255
    
    # Create a copy of the original image
    result = original_image.copy()
    
    # Set non-spot areas to white
    result[mask == 0] = [255, 255, 255]
    
    return result

# Function to segment dark spots based on RGB ranges
def segment_dark_spots(image_array):
    # Convert the image to RGB if it's not already
    if len(image_array.shape) == 2 or image_array.shape[2] == 1:
        image_rgb = cv2.cvtColor(image_array, cv2.COLOR_GRAY2RGB)
    else:
        image_rgb = cv2.cvtColor(image_array, cv2.COLOR_BGR2RGB)
    
    # Define RGB ranges for dark spots
    rgb_ranges = {
        "light_brown": ([139, 69, 19], [165, 105, 45]),
        "dark_brown": ([75, 45, 10], [110, 75, 35]),
        "black": ([0, 0, 0], [30, 30, 30]),
        "gray_green": ([100, 120, 100], [150, 160, 140]),
        "yellow_brown": ([150, 120, 20], [200, 160, 60]),
        "dark_purple": ([75, 0, 50], [120, 40, 80]),
        "speckled_brown": ([120, 80, 50], [150, 110, 80])
    }
    
    # Create a mask for all dark spots
    final_mask = np.zeros(image_rgb.shape[:2], dtype=np.uint8)
    
    # Allow user to adjust the sensitivity
    for color_name, (lower, upper) in rgb_ranges.items():
        # Create a mask for the current RGB range
        lower_bound = np.array(lower, dtype=np.uint8)
        upper_bound = np.array(upper, dtype=np.uint8)
        mask = cv2.inRange(image_rgb, lower_bound, upper_bound)
        
        # Add this mask to the final mask
        final_mask = cv2.bitwise_or(final_mask, mask)
    
    # Apply the mask to the original image
    segmented_image = cv2.bitwise_and(image_rgb, image_rgb, mask=final_mask)
    
    # Also create a white background version
    white_bg_image = create_white_bg_spots(final_mask, image_rgb)
    
    # Analyze the segmented spots and get spot information
    spot_info = analyze_spots(image_rgb, final_mask)
    
    return segmented_image, white_bg_image, final_mask, spot_info

# Function to analyze segmented spots and extract information
def analyze_spots(original_image, mask):
    # Find connected components in the mask
    num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(mask, connectivity=8)
    
    # Skip the first component (label 0) as it's the background
    spot_info = []
    total_image_pixels = original_image.shape[0] * original_image.shape[1]
    
    for i in range(1, num_labels):
        # Get the area of the spot
        area = stats[i, cv2.CC_STAT_AREA]
        
        # Skip very small spots (likely noise)
        if area < 10:
            continue
            
        # Create a mask for this specific spot
        spot_mask = (labels == i).astype(np.uint8)
        
        # Get the mean color of the spot
        mean_color = cv2.mean(original_image, mask=spot_mask)
        mean_color_rgb = (int(mean_color[0]), int(mean_color[1]), int(mean_color[2]))
        
        # Calculate the percentage of the image covered by this spot
        spot_percentage = (area / total_image_pixels) * 100
        
        # Calculate the bounding box
        x = stats[i, cv2.CC_STAT_LEFT]
        y = stats[i, cv2.CC_STAT_TOP]
        w = stats[i, cv2.CC_STAT_WIDTH]
        h = stats[i, cv2.CC_STAT_HEIGHT]
        
        # Add the spot information to the list
        spot_info.append({
            'spot_number': i,
            'size_pixels': area,
            'color_rgb': mean_color_rgb,
            'coverage_percentage': spot_percentage,
            'bbox': (x, y, w, h),
            'centroid': (int(centroids[i][0]), int(centroids[i][1]))
        })
    
    return spot_info

# Function to convert PIL Image to OpenCV format
def pil_to_cv2(pil_image):
    return cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)

# Function to convert OpenCV image to PIL
def cv2_to_pil(cv2_image):
    return Image.fromarray(cv2.cvtColor(cv2_image, cv2.COLOR_BGR2RGB))

# Function to draw bounding boxes and labels on image
def draw_spot_annotations(image, spot_info):
    # Create a copy of the image
    annotated_image = image.copy()
    
    # Draw bounding boxes and spot numbers
    for spot in spot_info:
        x, y, w, h = spot['bbox']
        cv2.rectangle(annotated_image, (x, y), (x+w, y+h), (0, 255, 0), 2)
        cv2.putText(annotated_image, f"#{spot['spot_number']}", 
                   (x, y-5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
    
    return annotated_image

# File uploader
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

# Sidebar for settings
st.sidebar.header("Settings")
remove_bg = st.sidebar.checkbox("Remove Background", value=True)
min_spot_size = st.sidebar.slider("Minimum Spot Size (pixels)", 5, 100, 10)
show_annotations = st.sidebar.checkbox("Show Spot Annotations", value=True)

# Process the image if it's uploaded
if uploaded_file is not None:
    # Load the image
    image = Image.open(uploaded_file)
    
    # Create two columns for the images
    col1, col2 = st.columns(2)
    
    with col1:
        st.subheader("Original Image")
        st.image(image, use_container_width=True)
    
    # Process the image
    if remove_bg:
        with st.spinner('Removing background...'):
            no_bg_img = remove_background(image)
        
        with col2:
            st.subheader("Background Removed")
            st.image(no_bg_img, use_container_width=True)
        
        # Convert PIL image to OpenCV format for further processing
        img_for_processing = pil_to_cv2(no_bg_img)
    else:
        img_for_processing = pil_to_cv2(image)
    
    # Segment dark spots
    with st.spinner('Analyzing spots...'):
        segmented_img, white_bg_spots, mask, spot_info = segment_dark_spots(img_for_processing)
    
    # Filter spots by minimum size
    filtered_spot_info = [spot for spot in spot_info if spot['size_pixels'] >= min_spot_size]
    
    # Recalculate total coverage
    total_coverage = sum(spot['coverage_percentage'] for spot in filtered_spot_info)
    
    # Create two more columns for the segmented images
    col3, col4 = st.columns(2)
    
    with col3:
        st.subheader("Detected Dark Spots")
        st.image(segmented_img, use_container_width=True)
    
    with col4:
        st.subheader("Spots on White Background")
        if show_annotations and filtered_spot_info:
            annotated_img = draw_spot_annotations(white_bg_spots, filtered_spot_info)
            st.image(annotated_img, use_container_width=True)
        else:
            st.image(white_bg_spots, use_container_width=True)
    
    # Display results
    st.subheader("Analysis Results")
    
    col_metrics1, col_metrics2, col_metrics3 = st.columns(3)
    
    with col_metrics1:
        st.metric("Total Spots", len(filtered_spot_info))
    
    with col_metrics2:
        st.metric("Total Coverage", f"{total_coverage:.2f}%")
    
    with col_metrics3:
        if len(filtered_spot_info) > 0:
            largest_spot = max(filtered_spot_info, key=lambda x: x['size_pixels'])
            st.metric("Largest Spot Size", f"{largest_spot['size_pixels']} pixels")
        else:
            st.metric("Largest Spot Size", "0 pixels")
    
    # Display detailed spot information
    if filtered_spot_info:
        st.subheader("Spot Details")
        
        # Convert spot info to a format suitable for a dataframe
        spot_data = []
        for spot in filtered_spot_info:
            color_hex = "#{:02x}{:02x}{:02x}".format(*spot['color_rgb'])
            spot_data.append({
                "Spot #": spot['spot_number'],
                "Size (pixels)": spot['size_pixels'],
                "Color": color_hex,
                "Coverage (%)": f"{spot['coverage_percentage']:.2f}%"
            })
        
        # Create tabs for different views
        tab1, tab2 = st.tabs(["Table View", "Detailed View"])
        
        with tab1:
            # Display as a table
            st.dataframe(spot_data)
        
        with tab2:
            # Display detailed information for each spot
            for spot in filtered_spot_info:
                with st.expander(f"Spot #{spot['spot_number']} - {spot['size_pixels']} pixels"):
                    col_a, col_b = st.columns([1, 3])
                    
                    with col_a:
                        # Show the color
                        color_hex = "#{:02x}{:02x}{:02x}".format(*spot['color_rgb'])
                        st.markdown(f"<div style='background-color: {color_hex}; width: 50px; height: 50px; border-radius: 5px;'></div>", unsafe_allow_html=True)
                    
                    with col_b:
                        st.write(f"Size: {spot['size_pixels']} pixels")
                        st.write(f"Coverage: {spot['coverage_percentage']:.2f}%")
                        st.write(f"RGB Color: {spot['color_rgb']}")
                        st.write(f"Position: {spot['centroid']}")
    else:
        st.info("No spots detected with the current settings. Try adjusting the minimum spot size.")

# Add some information about how to use the app
with st.expander("How to use this app"):
    st.write("""
    1. Upload an image of a fruit with visible dark spots or blemishes.
    2. The app will automatically remove the background (if selected) and detect dark spots.
    3. Adjust the minimum spot size to filter out noise or small spots.
    4. Toggle spot annotations to see numbered bounding boxes around the detected spots.
    5. View detailed information about each spot in the 'Spot Details' section.
    
    This app can help in:
    - Assessing fruit quality
    - Tracking disease progression
    - Quantifying surface blemishes
    """)

# Footer
st.markdown("---")
st.markdown("<p style='text-align: center;'>Developed by Manith Marapperuma</p>", unsafe_allow_html=True)