app.py

import streamlit as st
import requests
import re
import cv2
import numpy as np
from pathlib import Path
import tempfile
import time
from ultralytics import YOLO
import torch
import xml.etree.ElementTree as ET  # For loading water polygon
from shapely.geometry import Point, Polygon  # For water polygon logic
import os
import smtplib                 # For sending emails
from email.mime.text import MIMEText # For formatting emails

# ─── Streamlit page config & CSS ───────────────────────────────
st.set_page_config(page_title="Swimmer Detection", page_icon="🏊‍♂️", layout="centered")
st.markdown("""
<style>
.stButton>button { width:100%; background:#f0f2f6; color:#262730; border:none;
  padding:10px; border-radius:5px; margin:5px 0; }
.stSlider > div[data-baseweb="slider"] { color: #f63366; }
.big-title { font-size: 2.5rem; font-weight: 700; margin-bottom: 1rem; }
.section-header { font-size: 1.3rem; font-weight: 600; margin-top: 2rem; margin-bottom: 0.5rem; }
.result-count { font-size: 1.2rem; font-weight: 500; color: #f63366; }
</style>
""", unsafe_allow_html=True)

# ─── Session state defaults ─────────────────────────────────────
defaults = {
    'page': 1,
    'video_file': None,
    'video_url': "",
    'analysis_complete': False,
    'people_in_water_count': 0,
    'people_on_beach_count': 0,
    'processed_video_path': None
}
for key, val in defaults.items():
    if key not in st.session_state:
        st.session_state[key] = val

def navigate_to(page: int):
    st.session_state.page = max(1, min(3, page))

# ─── PAGE 1: Upload or Paste URL (with .m3u8 scraping) ──────────
def show_page_1():
    st.markdown('<div class="big-title">🏊‍♂️ Swimmer Detection App</div>', unsafe_allow_html=True)
    st.markdown('<div class="section-header">1️⃣ Upload or Paste Video URL</div>', unsafe_allow_html=True)
    st.write("Upload a local video or paste a public video/stream URL (e.g. Surfline link). For best results, use a short video.")

    # 1. capture raw input
    raw_url = st.text_input("Public video/stream URL:", st.session_state.video_url).strip()

    # 2. if it's not a direct .mp4/.mov/.avi/.m3u8, try to scrape an HLS playlist
    final_url = ""
    if raw_url:
        if raw_url.lower().endswith(('.mp4', '.mov', '.avi', '.m3u8')):
            final_url = raw_url
        else:
            st.info("Fetching page and looking for an HLS (`.m3u8`) playlist…")
            try:
                r = requests.get(raw_url, timeout=10)
                m = re.search(r'https://[^\s"\'<>]+?\.m3u8', r.text)
                if m:
                    final_url = m.group(0)
                    st.success("Found stream: " + final_url)
                else:
                    st.error("No `.m3u8` link found on that page.")
            except Exception as e:
                st.error(f"Failed to fetch page: {e}")

    # 3. save into session
    st.session_state.video_url = final_url

    # 4. preview side-by-side
    col1, col2 = st.columns(2)
    with col1:
        if final_url:
            st.video(final_url)
    with col2:
        uploaded = st.file_uploader("Or upload a 5-minute video:", type=['mp4','mov','avi'])
        if uploaded:
            st.session_state.video_file = uploaded
            st.session_state.video_url = ""  # clear any scraped URL
            st.video(uploaded)

    # 5. proceed
    if st.button("➡️ Analyse"):
        if not (st.session_state.video_file or st.session_state.video_url):
            st.warning("Please upload a file or supply a working video/stream URL.")
        else:
            st.session_state.analysis_complete = False
            navigate_to(2)

# ─── PAGE 2: Analysis In Progress ────────────────────────────────
MODEL_PATH           = "best.pt"
ANNOTATIONS_XML_PATH = "annotations.xml"
PERSON_CLASS_ID      = 0

@st.cache_resource
def load_model(path):
    try:
        model = YOLO(path)
        if torch.cuda.is_available():
            model.to("cuda")
        return model
    except Exception as e:
        st.error(f"Error loading model '{path}': {e}")
        return None

@st.cache_data
def load_polygon(xml_path, label="water"):
    p = Path(xml_path)
    if not p.exists():
        st.error(f"Annotation XML not found: {xml_path}")
        return None
    tree = ET.parse(p)
    root = tree.getroot()
    pts = []
    for poly in root.findall(".//polygon"):
        if poly.attrib.get("label") == label:
            for pt in poly.attrib["points"].split(';'):
                x, y = map(float, pt.split(','))
                pts.append((x, y))
            break
    if not pts:
        st.error(f"No polygon with label='{label}' in XML")
        return None
    return Polygon(pts)

def load_polygons(xml_path):
    water_poly = load_polygon(xml_path, label="water")
    flags_poly = load_polygon(xml_path, label="flags")
    return water_poly, flags_poly

# --- Email Notification Function ---
def send_email_notification(subject, body_text, recipient_email, total_water, total_beach):
    sender_email = "johnkim6400@gmail.com"
    sender_app_password = "zlhqkesmfrtdclkp"
    smtp_server_address = "smtp.gmail.com"
    smtp_port_str = "587"

    if not recipient_email:
        st.warning("No email address provided for notifications.")
        return False

    try:
        # Create email body with HTML formatting
        html_content = f"""
        <html>
            <body style="font-family: Arial, sans-serif;">
                <h2>Video Analysis Results</h2>
                <p>Your video analysis has been completed. Here are the results:</p>
                
                <div style="background-color: #f5f5f5; padding: 15px; border-radius: 5px; margin: 10px 0;">
                    <p><strong>People detected in water:</strong> {total_water}</p>
                    <p><strong>People detected on beach:</strong> {total_beach}</p>
                    <p><strong>Total people detected:</strong> {total_water + total_beach}</p>
                </div>
                
                <p>{body_text}</p>
                
                <p style="color: #666;">This is an automated notification from the AI Rip Current Checker system.</p>
            </body>
        </html>
        """

        msg = MIMEText(html_content, 'html')
        msg['Subject'] = subject
        msg['From'] = sender_email
        msg['To'] = recipient_email

        with smtplib.SMTP(smtp_server_address, int(smtp_port_str)) as server:
            server.ehlo()
            server.starttls()
            server.ehlo()
            server.login(sender_email, sender_app_password)
            server.sendmail(sender_email, recipient_email, msg.as_string())
            st.success(f"Analysis results sent to {recipient_email}")
            return True

    except Exception as e:
        st.error(f"Failed to send email notification: {e}")
        return False

def process_video(source, model, water_poly, max_frames_to_process=None):
    cap = cv2.VideoCapture(source)
    if not cap.isOpened():
        st.error(f"Cannot open video/stream: {source}")
        return None, 0, 0

    w   = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    h   = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap.get(cv2.CAP_PROP_FPS) or 30
    if max_frames_to_process is not None:
        total = min(int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) or -1, max_frames_to_process)
    else:
        total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) or -1

    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
    tmp.close()
    out = cv2.VideoWriter(tmp.name,
                          cv2.VideoWriter_fourcc(*'mp4v'),
                          fps, (w, h))

    water_count = beach_count = 0
    progress = st.progress(0.0)
    status = st.empty()
    frame_i = 0

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        results = model(frame, verbose=False)[0]
        for box in results.boxes:
            if int(box.cls[0]) != PERSON_CLASS_ID:
                continue
            x1, y1, x2, y2 = map(int, box.xyxy[0])
            cx, cy = (x1 + x2)//2, y2
            in_water = water_poly.contains(Point(cx, cy))
            color = (255, 0, 0) if in_water else (0, 255, 0)
            label = "Water" if in_water else "Beach"
            if in_water:
                water_count += 1
            else:
                beach_count += 1

            cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
            cv2.putText(frame, label, (x1, y1-10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)

        text = f"In Water: {water_count} | On Beach: {beach_count}"
        cv2.putText(frame, text, (10, 30),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2)

        out.write(frame)
        frame_i += 1

        if total > 0:
            progress.progress(frame_i/total)
            status.text(f"Frame {frame_i}/{total}")
        else:
            status.text(f"Frame {frame_i}")

        if max_frames_to_process and frame_i >= max_frames_to_process:
            break

    cap.release()
    out.release()
    return tmp.name, water_count, beach_count

def show_page_2():
    st.markdown('<div class="section-header">2️⃣ Live Analysis</div>', unsafe_allow_html=True)
    
    # Initialize variables at the function level
    cap = None
    out_writer = None
    processed_video_temp_file = None
    source_for_processing = None
    
    try:
        col_top1, col_top2 = st.columns([1, 8])
        with col_top1:
            if st.button("⏪ Back to Upload", key="back_to_upload_top"):
                st.session_state.stop_analysis = True
                navigate_to(1)
                return

        src_file = st.session_state.get('video_file', None)
        src_url = st.session_state.get('video_url', "")
        use_cam = st.session_state.get('use_webcam', False)

        # --- Sidebar controls ---
        st.sidebar.markdown("## ⚙️ Analysis Controls")
        playback_speed = st.sidebar.select_slider(
            "Playback Speed",
            options=[0.25, 0.5, 1.0, 2.0, 4.0],
            value=1.0,
            help="Set to 1.0 for real-time, higher for faster playback."
        )

        alert_email = st.sidebar.text_input("Alert Email Address", value=st.session_state.get('alert_email', ''))
        st.session_state.alert_email = alert_email

        if 'stop_analysis' not in st.session_state:
            st.session_state.stop_analysis = False

        if st.button("🛑 Stop Analysis", key="stop_live_analysis"):
            st.session_state.stop_analysis = True
            st.info("Analysis stop requested. Processing will halt after the current frame.")

        # Determine video source
        if use_cam:
            source_for_processing = 0
            st.info("Attempting to use webcam...")
        elif src_file:
            try:
                with tempfile.NamedTemporaryFile(delete=False, suffix=Path(src_file.name).suffix) as tmp_video_file:
                    tmp_video_file.write(src_file.getbuffer())
                    source_for_processing = tmp_video_file.name
                st.info(f"Processing uploaded file: {src_file.name}")
            except Exception as e:
                st.error(f"Error with uploaded file: {e}")
                if st.button("⏪ Back to Upload", key="err_back_upload_file"):
                    navigate_to(1)
                return
        elif src_url:
            source_for_processing = src_url
            st.info(f"Processing URL: {src_url}")
        else:
            st.warning("No video source selected.")
            if st.button("⏪ Back to Upload", key="back_upload_no_source"):
                navigate_to(1)
            return

        notification_container = st.container()
        model = load_model(MODEL_PATH)
        water_poly, flags_poly = load_polygons(ANNOTATIONS_XML_PATH)

        # Placeholders for UI elements
        st.markdown("---")
        col_counts1, col_counts2 = st.columns(2)
        with col_counts1:
            water_count_placeholder = st.empty()
            beach_count_placeholder = st.empty()
        with col_counts2:
            inside_flags_placeholder = st.empty()
            outside_flags_placeholder = st.empty()
        total_detected_placeholder = st.empty()
        st.markdown("---")
        frame_placeholder = st.empty()

        # Open video capture
        cap = cv2.VideoCapture(source_for_processing)
        if not cap.isOpened():
            st.error(f"Cannot open video source: {source_for_processing}")
            return

        # Get video properties
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        fps = cap.get(cv2.CAP_PROP_FPS) or 30
        frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        frame_delay = 1.0 / (fps * playback_speed) if playback_speed > 0 else 0

        # --- Frame Control (moved here after we have valid video) ---
        st.sidebar.markdown("### 🎞️ Frame Control")
        col1, col2 = st.sidebar.columns(2)
        with col1:
            start_frame = st.number_input("Start Frame", 
                                        min_value=0, 
                                        max_value=max(0, total_frames-2),  # Ensure there's at least one frame to process
                                        value=0)
        with col2:
            end_frame = st.number_input("End Frame", 
                                      min_value=max(1, start_frame+1),  # Ensure end_frame is at least 1 and greater than start_frame
                                      max_value=total_frames, 
                                      value=min(total_frames, start_frame + 100))  # Default to processing 100 frames
        
        frame_skip = st.sidebar.slider("Process every Nth frame", 
                                     min_value=1, 
                                     max_value=30, 
                                     value=1, 
                                     help="Higher values = faster processing but might miss some actions")

        # Initialize video writer
        processed_video_temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
        out_writer = cv2.VideoWriter(
            processed_video_temp_file.name,
            cv2.VideoWriter_fourcc(*'mp4v'),
            fps,
            (frame_width, frame_height)
        )

        # Set initial frame position
        cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
        current_frame = start_frame
        total_frames_to_process = end_frame - start_frame

        # Initialize lists to track counts
        water_counts = []
        beach_counts = []

        # Progress tracking
        progress_bar = st.progress(0)
        status_text = st.empty()

        while True:
            if st.session_state.stop_analysis:
                st.warning("Analysis stopped by user.")
                break

            ret, frame = cap.read()
            if not ret or current_frame >= end_frame:
                st.write("Debug: Reached end of video or failed to read frame")
                break

            # Skip frames according to frame_skip setting
            if (current_frame - start_frame) % frame_skip != 0:
                current_frame += 1
                continue

            # Update progress
            progress = (current_frame - start_frame) / total_frames_to_process
            progress_bar.progress(progress)
            status_text.text(f"Processing frame {current_frame} of {end_frame} (every {frame_skip} frames)")

            # Process frame with model
            results = model(frame, verbose=False)[0]
            
            # Reset frame counts
            wc, bc, ifc, ofc = 0, 0, 0, 0

            # Process detections
            for box in results.boxes:
                detected_class_id = int(box.cls[0])
                if detected_class_id != PERSON_CLASS_ID:
                    continue

                x1, y1, x2, y2 = map(int, box.xyxy[0])
                cx, cy = (x1 + x2) // 2, y2

                in_water = water_poly.contains(Point(cx, cy)) if water_poly else False
                in_flags = flags_poly.contains(Point(cx, cy)) if flags_poly else False

                color = (255, 0, 0) if in_water else (0, 255, 0)
                label = "Water" if in_water else "Beach"

                if in_water:
                    wc += 1
                else:
                    bc += 1

                if in_flags:
                    ifc += 1
                else:
                    ofc += 1

                # Draw bounding box and label
                cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
                cv2.putText(frame, label, (x1, y1-10),
                          cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)

            # Add counts to tracking lists
            water_counts.append(wc)
            beach_counts.append(bc)

            # Add counts to frame
            cv2.putText(frame, f"In Water: {wc} | On Beach: {bc}", (10, 30),
                      cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2)

            # Update UI
            frame_placeholder.image(frame, channels="BGR", use_container_width=True)
            water_count_placeholder.markdown(f'<div class="result-count">In Water: <b>{wc}</b></div>', unsafe_allow_html=True)
            beach_count_placeholder.markdown(f'<div class="result-count">On Beach: <b>{bc}</b></div>', unsafe_allow_html=True)
            total_detected_placeholder.markdown(f'<div class="result-count">Total People Detected: <b>{wc + bc}</b></div>', unsafe_allow_html=True)

            # Write frame to output video
            out_writer.write(frame)

            # Add frame delay back
            if frame_delay > 0:
                time.sleep(frame_delay)

            current_frame += 1

        # After the video processing loop, before the finally block, add:
        if not st.session_state.stop_analysis and alert_email:
            # Calculate final counts (maximum number of people detected in any frame)
            total_water_count = max(water_counts) if water_counts else 0
            total_beach_count = max(beach_counts) if beach_counts else 0
            
            # Send email notification
            notification_text = "You can view the full analysis and processed video in the application."
            send_email_notification(
                subject="Video Analysis Complete - AI Rip Current Checker",
                body_text=notification_text,
                recipient_email=alert_email,
                total_water=total_water_count,
                total_beach=total_beach_count
            )

    except Exception as e:
        st.error(f"An error occurred: {str(e)}")
        import traceback
        st.error(traceback.format_exc())
        
    finally:
        # Cleanup
        if cap is not None:
            cap.release()
        if out_writer is not None:
            out_writer.release()

        # Save and display the processed video
        if processed_video_temp_file is not None and os.path.exists(processed_video_temp_file.name) and os.path.getsize(processed_video_temp_file.name) > 0:
            st.session_state.processed_video_path = processed_video_temp_file.name
            st.success("✅ Analysis complete. Displaying processed video below:")
            st.video(processed_video_temp_file.name)
        else:
            st.error("Failed to save processed video.")

        # Clean up temporary files
        if src_file and source_for_processing and os.path.exists(source_for_processing):
            try:
                os.unlink(source_for_processing)
            except Exception as e:
                st.warning(f"Could not delete temporary file: {e}")

    # --- Back and review controls at the bottom ---
    st.markdown("---")
    col1, col2 = st.columns(2)
    with col1:
        if st.button("⏪ Back to Upload", key="back_to_upload_bottom"):
            navigate_to(1)
    with col2:
        if st.button("🔄 Analyse Another Video", key="analyse_another_live"):
            for k in ['video_file','video_url','analysis_complete']:
                st.session_state[k] = defaults[k]
            navigate_to(1)

# ─── PAGE 3: Results ─────────────────────────────────────────────
def show_page_3():
    st.title("3️⃣ Results")
    if not st.session_state.analysis_complete:
        st.warning("No results to show yet.")
        if st.button("⏪ Back"): navigate_to(2)
        return

    st.subheader("Annotated Video")
    st.video(st.session_state.processed_video_path)

    st.subheader("Final Counts")
    st.markdown(f"- **In Water:** {st.session_state.people_in_water_count}")
    st.markdown(f"- **On Beach:** {st.session_state.people_on_beach_count}")
    st.markdown(f"- **Total:** {st.session_state.people_in_water_count + st.session_state.people_on_beach_count}")

    if st.button("🔄 Analyse Again"):
        for k in ['video_file','video_url','analysis_complete']:
            st.session_state[k] = defaults[k]
        navigate_to(1)

# ─── App entrypoint ─────────────────────────────────────────────
def main():
    if st.session_state.page == 1:
        show_page_1()
    elif st.session_state.page == 2:
        show_page_2()
    else:
        show_page_3()

if __name__ == "__main__":
    main()