import gradio as gr
import cv2
import numpy as np
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Flatten
from PIL import Image, ImageDraw
import random
import time

# ------------------------------
# Dummy CNN model
# ------------------------------
def get_dummy_model():
    model = Sequential([
        Flatten(input_shape=(64,64,3)),
        Dense(2, activation='softmax')
    ])
    return model

model = get_dummy_model()

# ------------------------------
# Load sample video
# ------------------------------
video_path = "sample_video.mp4"
cap = cv2.VideoCapture(video_path)

# ------------------------------
# State variables
# ------------------------------
frame_counter = 0
obstacle_count = 0

# ------------------------------
# Function to get next frame
# ------------------------------
def get_next_frame(speed_slider):
    global frame_counter, obstacle_count, cap
    ret, frame = cap.read()
    if not ret:
        cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
        frame_counter = 0
        ret, frame = cap.read()
    frame_counter += 1

    # Resize for dummy CNN
    img = cv2.resize(frame, (64,64))
    img_norm = img / 255.0
    img_input = np.expand_dims(img_norm, axis=0)

    # Dummy obstacle detection
    label = random.choice([0,1])  # 0=clear, 1=obstacle

    # Overlay rectangle if obstacle detected
    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    frame_pil = Image.fromarray(frame_rgb)
    draw = ImageDraw.Draw(frame_pil)
    status_text = "✅ Path Clear"

    if label == 1:
        status_text = f"⚠️ Obstacle Detected! (speed={speed_slider})"
        obstacle_count += 1
        # Draw random rectangle for demo
        x0, y0 = random.randint(50,200), random.randint(50,200)
        x1, y1 = x0+50, y0+50
        draw.rectangle([x0,y0,x1,y1], outline="red", width=4)

    # Display obstacle count
    status_text += f" | Total Obstacles: {obstacle_count}"

    return frame_pil, status_text

# ------------------------------
# Build Gradio Interface
# ------------------------------
with gr.Blocks() as demo:
    gr.Markdown("## 🛸 Autonomous Drone Demo (Enhanced Gradio Version)")
    gr.Markdown("Demo: AI-based obstacle detection with auto-play, overlay rectangles, speed slider, and obstacle counter.")

    speed_slider = gr.Slider(10, 50, value=20, label="Drone Speed (Demo)")
    output_img = gr.Image()
    output_text = gr.Textbox(label="Status")
    
    run_btn = gr.Button("Next Frame")
    
    run_btn.click(fn=get_next_frame, inputs=[speed_slider], outputs=[output_img, output_text])

demo.launch()