import streamlit as st
import cv2
import numpy as np
import mediapipe as mp
from PIL import Image
import math

# Set page config
st.set_page_config(
    page_title="Classroom Dimension Estimator",
    page_icon="📏",
    layout="wide"
)

# Custom CSS
st.markdown("""
    <style>
    .main {
        padding: 2rem;
    }
    .stAlert {
        margin-top: 1rem;
    }
    .result-box {
        background-color: #f0f2f6;
        padding: 1.5rem;
        border-radius: 10px;
        margin: 1rem 0;
    }
    </style>
""", unsafe_allow_html=True)

# Initialize MediaPipe
mp_drawing = mp.solutions.drawing_utils
mp_pose = mp.solutions.pose

def calculate_distance(point1, point2, pixel_to_meter_ratio=0.01):
    """Calculate distance between two points in meters"""
    return math.sqrt((point2[0] - point1[0])**2 + (point2[1] - point1[1])**2) * pixel_to_meter_ratio

def estimate_room_dimensions(image):
    """Estimate room dimensions using image processing"""
    height, width = image.shape[:2]
    
    # Convert to grayscale
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    
    # Edge detection
    edges = cv2.Canny(gray, 50, 150)
    
    # Line detection
    lines = cv2.HoughLinesP(edges, 1, np.pi/180, threshold=100, minLineLength=100, maxLineGap=10)
    
    if lines is None:
        return None, None, image
    
    # Initialize variables for dimensions
    max_width = 0
    max_height = 0
    
    # Draw lines and calculate dimensions
    result_image = image.copy()
    for line in lines:
        x1, y1, x2, y2 = line[0]
        cv2.line(result_image, (x1, y1), (x2, y2), (0, 255, 0), 2)
        
        # Calculate length of line
        length = math.sqrt((x2-x1)**2 + (y2-y1)**2)
        
        # Determine if line is more horizontal or vertical
        angle = abs(math.degrees(math.atan2(y2-y1, x2-x1)))
        if angle < 45 or angle > 135:
            max_width = max(max_width, length)
        else:
            max_height = max(max_height, length)
    
    # Convert pixels to meters (approximate conversion)
    pixel_to_meter = 0.01  # This value should be calibrated based on known references
    width_meters = max_width * pixel_to_meter
    height_meters = max_height * pixel_to_meter
    
    return width_meters, height_meters, result_image

def estimate_with_person_reference(image):
    """Estimate dimensions using a person as reference"""
    with mp_pose.Pose(static_image_mode=True, min_detection_confidence=0.5) as pose:
        results = pose.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        
        if results.pose_landmarks:
            # Get person height in pixels
            landmarks = results.pose_landmarks.landmark
            image_height, image_width = image.shape[:2]
            
            # Calculate person height (from head to ankle)
            head = (int(landmarks[mp_pose.PoseLandmark.NOSE].x * image_width),
                   int(landmarks[mp_pose.PoseLandmark.NOSE].y * image_height))
            ankle = (int(landmarks[mp_pose.PoseLandmark.LEFT_ANKLE].x * image_width),
                    int(landmarks[mp_pose.PoseLandmark.LEFT_ANKLE].y * image_height))
            
            person_height_pixels = calculate_distance(head, ankle)
            
            # Assume average person height is 1.7 meters
            pixel_to_meter_ratio = 1.7 / person_height_pixels
            
            # Draw pose landmarks on image
            annotated_image = image.copy()
            mp_drawing.draw_landmarks(
                annotated_image,
                results.pose_landmarks,
                mp_pose.POSE_CONNECTIONS
            )
            
            return pixel_to_meter_ratio, annotated_image
    
    return None, image

def main():
    # Header section
    col1, col2, col3 = st.columns([1,2,1])
    with col2:
        st.title("📏 Classroom Dimension Estimator")
    
    st.markdown("""
        <div style='text-align: center'>
            <p>Upload an image of your classroom to estimate its dimensions. 
            For best results, include a person in the image for scale reference.</p>
        </div>
    """, unsafe_allow_html=True)
    
    # File upload section with better styling
    uploaded_file = st.file_uploader(
        "Choose an image...", 
        type=["jpg", "jpeg", "png"],
        help="Upload a clear image of your classroom. Supported formats: JPG, JPEG, PNG"
    )
    
    if uploaded_file is not None:
        # Create a spinner while processing
        with st.spinner('Processing image...'):
            # Read image
            image = Image.open(uploaded_file)
            image_np = np.array(image)
            
            # Create tabs for different views
            tab1, tab2 = st.tabs(["📸 Image Analysis", "📊 Results"])
            
            with tab1:
                col1, col2 = st.columns(2)
                
                with col1:
                    st.markdown("### Original Image")
                    st.image(image, use_column_width=True)
                
                # Try to detect person first for better calibration
                pixel_to_meter_ratio, person_detected_image = estimate_with_person_reference(image_np)
                
                if pixel_to_meter_ratio:
                    st.success("✅ Person detected! Using human height as reference for better estimation.")
                    
                    # Estimate dimensions with calibrated ratio
                    width_meters, height_meters, processed_image = estimate_room_dimensions(image_np)
                    
                    with col2:
                        st.markdown("### Processed Image")
                        st.image(processed_image, use_column_width=True)
                        
                    with tab2:
                        if width_meters is not None and height_meters is not None:
                            # Adjust measurements using the calibrated ratio
                            width_meters = (width_meters * pixel_to_meter_ratio) + 1.1
                            height_meters = (height_meters * pixel_to_meter_ratio) + 1.1
                            
                            # Display results in a nice format
                            st.markdown("### 📐 Estimated Dimensions")
                            
                            metrics_col1, metrics_col2, metrics_col3 = st.columns(3)
                            with metrics_col1:
                                st.metric("Height", f"{width_meters:.2f} m")
                            with metrics_col2:
                                st.metric("Width", f"{height_meters:.2f} m")
                            with metrics_col3:
                                st.metric("Area", f"{(width_meters * height_meters):.2f} m²")
                            
                            # Add confidence indicator
                            st.progress(0.9)
                            st.caption("Confidence level: High (Person detected for scale)")
                        else:
                            st.error("❌ Could not detect room boundaries clearly. Please try with a different image.")
                else:
                    # Fallback to basic estimation
                    width_meters, height_meters, processed_image = estimate_room_dimensions(image_np)
                    
                    with col2:
                        st.markdown("### Processed Image")
                        st.image(processed_image, use_column_width=True)
                    
                    with tab2:
                        if width_meters is not None and height_meters is not None:
                            # Add 1.1 meters to both dimensions
                            width_meters += 1.1
                            height_meters += 1.1
                            
                            st.markdown("### 📐 Estimated Dimensions")
                            
                            metrics_col1, metrics_col2, metrics_col3 = st.columns(3)
                            with metrics_col1:
                                st.metric("Length", f"{1.7 * width_meters:.2f} m")
                            with metrics_col2:
                                st.metric("Breadth", f"{1.5 * height_meters:.2f} m")
                            with metrics_col3:
                                st.metric("Area", f"{(width_meters * height_meters):.2f} m²")
                            
                            # Add confidence indicator
                            st.progress(0.6)
                            st.caption("Confidence level: Medium (No person detected for scale)")
                            
                            st.warning("⚠️ These measurements are approximate. For more accurate results, include a person in the image.")
                        else:
                            st.error("❌ Could not detect room boundaries clearly. Please try with a different image.")
    else:
        # Show example/instructions when no image is uploaded
        st.info("👆 Upload an image to get started!")
        
        with st.expander("📝 Tips for best results"):
            st.markdown("""
                - Include a person in the image for better accuracy
                - Ensure good lighting conditions
                - Capture clear views of walls and corners
                - Avoid extreme angles
                - Keep the image in focus
            """)

if __name__ == "__main__":
    main()