#!/usr/bin/env python
"""
Example of AI-powered features in FMUS-VID.

This script demonstrates object detection, tracking, and background removal.
"""

import sys
import os
import time
from pathlib import Path
import cv2
import numpy as np

# Add the parent directory to the path so we can import fmusvid
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from fmusvid import Video
from fmusvid.ai.detection import ObjectDetector
from fmusvid.ai.tracking import ObjectTracker
from fmusvid.ai.background import BackgroundRemover, BackgroundBlur

def run_object_detection(video_path):
    """Run object detection on a video file."""
    print(f"Running object detection on {video_path}")

    try:
        # Load video
        cap = cv2.VideoCapture(str(video_path))
        if not cap.isOpened():
            print(f"Error: Could not open video {video_path}")
            return None

        # Get video properties
        fps = cap.get(cv2.CAP_PROP_FPS)
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

        # Create output video
        output_dir = os.path.dirname(video_path)
        output_filename = f"{os.path.splitext(os.path.basename(video_path))[0]}_detected.mp4"
        output_path = os.path.join(output_dir, output_filename)

        out = cv2.VideoWriter(
            str(output_path),
            cv2.VideoWriter_fourcc(*'mp4v'),
            fps,
            (width, height)
        )

        # Initialize object detector
        try:
            detector = ObjectDetector()
            print(f"Loaded {detector.model_name} model on {detector.device}")
            has_detector = True
        except Exception as e:
            print(f"Error initializing object detector: {e}")
            print("Falling back to drawing placeholders instead")
            has_detector = False

        # Process video
        frame_count = 0
        start_time = time.time()

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

            # Run detection or create placeholder visualization
            if has_detector:
                try:
                    # Detect objects
                    detections = detector.detect(frame, conf_threshold=0.5)

                    # Annotate frame
                    result = detector.annotate_frame(frame, detections)

                    # Add detection count
                    cv2.putText(
                        result,
                        f"Detected: {len(detections)} objects",
                        (10, 30),
                        cv2.FONT_HERSHEY_SIMPLEX,
                        1.0,
                        (0, 255, 0),
                        2
                    )
                except Exception as e:
                    print(f"Error running detection: {e}")
                    result = frame.copy()  # Use original frame if detection fails

                    # Add error message
                    cv2.putText(
                        result,
                        "Detection error",
                        (10, 30),
                        cv2.FONT_HERSHEY_SIMPLEX,
                        1.0,
                        (0, 0, 255),
                        2
                    )
            else:
                # Create a placeholder visualization
                result = frame.copy()

                # Draw some random detection boxes as placeholders
                for i in range(3):
                    x = np.random.randint(10, width - 100)
                    y = np.random.randint(10, height - 100)
                    w = np.random.randint(50, 200)
                    h = np.random.randint(50, 200)

                    # Ensure box is within frame
                    if x + w >= width: w = width - x - 1
                    if y + h >= height: h = height - y - 1

                    # Draw box
                    cv2.rectangle(result, (x, y), (x + w, y + h), (0, 255, 0), 2)

                    # Draw label
                    labels = ["Person", "Car", "Dog", "Cat", "Bird"]
                    label = labels[i % len(labels)]
                    cv2.putText(
                        result,
                        label,
                        (x, y - 10),
                        cv2.FONT_HERSHEY_SIMPLEX,
                        0.5,
                        (0, 255, 0),
                        1
                    )

                # Add simulation message
                cv2.putText(
                    result,
                    "Simulated detection (no model available)",
                    (10, 30),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    0.7,
                    (0, 255, 255),
                    2
                )

            # Write frame
            out.write(result)

            # Update progress
            frame_count += 1
            if frame_count % 10 == 0:
                elapsed = time.time() - start_time
                print(f"Processed {frame_count} frames in {elapsed:.2f}s ({frame_count/elapsed:.2f} fps)")

        # Clean up
        cap.release()
        out.release()
        print(f"Saved detection results to {output_path}")
        return output_path

    except Exception as e:
        print(f"Error in object detection: {e}")
        return None

def run_object_tracking(video_path):
    """Run object tracking on a video file."""
    print(f"Running object tracking on {video_path}")

    try:
        # Load video
        cap = cv2.VideoCapture(str(video_path))
        if not cap.isOpened():
            print(f"Error: Could not open video {video_path}")
            return None

        # Get video properties
        fps = cap.get(cv2.CAP_PROP_FPS)
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

        # Create output video
        output_dir = os.path.dirname(video_path)
        output_filename = f"{os.path.splitext(os.path.basename(video_path))[0]}_tracked.mp4"
        output_path = os.path.join(output_dir, output_filename)

        out = cv2.VideoWriter(
            str(output_path),
            cv2.VideoWriter_fourcc(*'mp4v'),
            fps,
            (width, height)
        )

        # Initialize object detector and tracker
        try:
            detector = ObjectDetector()
            tracker = ObjectTracker()
            print(f"Loaded {detector.model_name} model on {detector.device}")
            has_tracking = True
        except Exception as e:
            print(f"Error initializing object detector or tracker: {e}")
            print("Falling back to drawing placeholders instead")
            has_tracking = False

        # Process video
        frame_count = 0
        start_time = time.time()

        # For placeholder tracking
        placeholder_tracks = []

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

            # Run tracking or create placeholder visualization
            if has_tracking:
                try:
                    # Only run detection every 5 frames for performance
                    if frame_count % 5 == 0:
                        detections = detector.detect(frame, conf_threshold=0.5)
                    else:
                        detections = []

                    # Update tracker
                    tracks = tracker.update(detections)

                    # Annotate frame
                    result = tracker.annotate_frame(frame)

                    # Add tracking count
                    cv2.putText(
                        result,
                        f"Tracking: {len(tracks)} objects",
                        (10, 30),
                        cv2.FONT_HERSHEY_SIMPLEX,
                        1.0,
                        (0, 255, 0),
                        2
                    )
                except Exception as e:
                    print(f"Error running tracking: {e}")
                    result = frame.copy()  # Use original frame if tracking fails

                    # Add error message
                    cv2.putText(
                        result,
                        "Tracking error",
                        (10, 30),
                        cv2.FONT_HERSHEY_SIMPLEX,
                        1.0,
                        (0, 0, 255),
                        2
                    )
            else:
                # Create a placeholder visualization with simulated tracking
                result = frame.copy()

                # Initialize placeholder tracks if needed
                if frame_count == 0 or len(placeholder_tracks) == 0:
                    placeholder_tracks = []
                    for i in range(3):
                        x = np.random.randint(10, width - 100)
                        y = np.random.randint(10, height - 100)
                        w = np.random.randint(50, 150)
                        h = np.random.randint(50, 150)
                        dx = np.random.randint(-5, 5)
                        dy = np.random.randint(-5, 5)

                        # Ensure box is within frame
                        if x + w >= width: w = width - x - 1
                        if y + h >= height: h = height - y - 1

                        # Create track
                        placeholder_tracks.append({
                            "id": i + 1,
                            "x": x,
                            "y": y,
                            "w": w,
                            "h": h,
                            "dx": dx,
                            "dy": dy,
                            "label": ["Person", "Car", "Dog"][i % 3]
                        })

                # Update and draw placeholder tracks
                for track in placeholder_tracks:
                    # Update position
                    track["x"] += track["dx"]
                    track["y"] += track["dy"]

                    # Bounce off edges
                    if track["x"] < 0 or track["x"] + track["w"] >= width:
                        track["dx"] *= -1
                        track["x"] = max(0, min(width - track["w"], track["x"]))

                    if track["y"] < 0 or track["y"] + track["h"] >= height:
                        track["dy"] *= -1
                        track["y"] = max(0, min(height - track["h"], track["y"]))

                    # Draw track
                    cv2.rectangle(
                        result,
                        (int(track["x"]), int(track["y"])),
                        (int(track["x"] + track["w"]), int(track["y"] + track["h"])),
                        (0, 255, 255),
                        2
                    )

                    # Draw ID and label
                    cv2.putText(
                        result,
                        f"ID: {track['id']} ({track['label']})",
                        (int(track["x"]), int(track["y"]) - 10),
                        cv2.FONT_HERSHEY_SIMPLEX,
                        0.5,
                        (0, 255, 255),
                        1
                    )

                # Add simulation message
                cv2.putText(
                    result,
                    "Simulated tracking (no model available)",
                    (10, 30),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    0.7,
                    (0, 255, 255),
                    2
                )

            # Write frame
            out.write(result)

            # Update progress
            frame_count += 1
            if frame_count % 10 == 0:
                elapsed = time.time() - start_time
                print(f"Processed {frame_count} frames in {elapsed:.2f}s ({frame_count/elapsed:.2f} fps)")

        # Clean up
        cap.release()
        out.release()
        print(f"Saved tracking results to {output_path}")
        return output_path

    except Exception as e:
        print(f"Error in object tracking: {e}")
        return None

def run_background_removal(video_path, background_path=None):
    """Run background removal or replacement on a video file."""
    print(f"Running background removal on {video_path}")

    try:
        # Load video
        cap = cv2.VideoCapture(str(video_path))
        if not cap.isOpened():
            print(f"Error: Could not open video {video_path}")
            return None

        # Get video properties
        fps = cap.get(cv2.CAP_PROP_FPS)
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

        # Load background if provided
        background_img = None
        if background_path and os.path.exists(background_path):
            try:
                background_img = cv2.imread(str(background_path))
                if background_img is not None:
                    background_img = cv2.cvtColor(background_img, cv2.COLOR_BGR2RGB)
                    background_img = cv2.resize(background_img, (width, height))
                    print(f"Loaded background image: {background_path}")
                else:
                    print(f"Failed to load background image: {background_path}")
                    background_img = None
            except Exception as e:
                print(f"Error loading background image: {e}")
                background_img = None

        # Determine output format
        output_dir = os.path.dirname(video_path)
        if background_img is not None:
            output_filename = f"{os.path.splitext(os.path.basename(video_path))[0]}_bg_replaced.mp4"
        else:
            output_filename = f"{os.path.splitext(os.path.basename(video_path))[0]}_bg_removed.mp4"

        output_path = os.path.join(output_dir, output_filename)
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(str(output_path), fourcc, fps, (width, height))

        # Initialize background remover
        try:
            if background_img is not None:
                bg_processor = BackgroundRemover()
            else:
                bg_processor = BackgroundBlur()
            print(f"Loaded {bg_processor.model_name if hasattr(bg_processor, 'model_name') else 'background blur'} model")
            has_bg_processor = True
        except Exception as e:
            print(f"Error initializing background remover: {e}")
            print("Falling back to simple background processing")
            has_bg_processor = False

        # Process video
        frame_count = 0
        start_time = time.time()

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

            # Process frame
            try:
                if has_bg_processor:
                    # Apply background removal/replacement
                    if background_img is not None:
                        # With custom background
                        result = bg_processor.process(frame, background=background_img)
                    else:
                        # With blur
                        result = bg_processor.process(frame)
                else:
                    # Create a simple effect if models are not available
                    result = frame.copy()

                    # Create a simple mask (circular gradient from center) - vectorized
                    y_coords, x_coords = np.meshgrid(np.arange(height), np.arange(width), indexing='ij')
                    center_x, center_y = width // 2, height // 2
                    dist = np.sqrt((x_coords - center_x)**2 + (y_coords - center_y)**2) / np.sqrt((width//2)**2 + (height//2)**2)
                    mask = np.maximum(0, 1 - dist).astype(np.float32)

                    # Apply simple foreground/background processing
                    if background_img is not None:
                        # Apply mask with background
                        mask_3d = np.repeat(mask[:, :, np.newaxis], 3, axis=2)
                        result = (frame * mask_3d + background_img * (1 - mask_3d)).astype(np.uint8)
                    else:
                        # Apply blur effect
                        blurred = cv2.GaussianBlur(frame, (25, 25), 0)
                        mask_3d = np.repeat(mask[:, :, np.newaxis], 3, axis=2)
                        result = (frame * mask_3d + blurred * (1 - mask_3d)).astype(np.uint8)

                    # Add simulation message
                    cv2.putText(
                        result,
                        "Simulated background effect (no model available)",
                        (10, 30),
                        cv2.FONT_HERSHEY_SIMPLEX,
                        0.7,
                        (0, 255, 255),
                        2
                    )
            except Exception as e:
                print(f"Error processing frame {frame_count}: {e}")
                result = frame.copy()

                # Add error message
                cv2.putText(
                    result,
                    "Processing error",
                    (10, 30),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    1.0,
                    (0, 0, 255),
                    2
                )

            # Write frame
            out.write(cv2.cvtColor(result, cv2.COLOR_RGB2BGR))

            # Update progress
            frame_count += 1
            if frame_count % 10 == 0:
                elapsed = time.time() - start_time
                print(f"Processed {frame_count} frames in {elapsed:.2f}s ({frame_count/elapsed:.2f} fps)")

        # Clean up
        cap.release()
        out.release()
        print(f"Saved background processing results to {output_path}")
        return output_path

    except Exception as e:
        print(f"Error in background removal: {e}")
        return None

def main():
    """Run the AI features example."""
    print("FMUS-VID AI Features Example")
    print("============================")

    # Hardcoded sample videos to try
    sample_video_paths = [
        os.path.join("test_media", "test_video.mp4"),
        os.path.join("test_media", "sample.mp4"),
        os.path.join("test_media", "sample2.mp4")
    ]

    # Find the first existing sample video
    video_path = None
    for path in sample_video_paths:
        if os.path.exists(path):
            video_path = path
            break

    # If no sample videos found, create a simple test video
    if video_path is None:
        print("No sample videos found. Creating a simple test video...")

        # Create test_media directory if it doesn't exist
        os.makedirs("test_media", exist_ok=True)

        # Create a simple test video with colored rectangles moving
        test_video_path = os.path.join("test_media", "generated_test.mp4")

        # Video properties
        width, height = 640, 480
        fps = 30
        duration = 5  # seconds

        # Create video writer
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(test_video_path, fourcc, fps, (width, height))

        # Generate frames with moving rectangles
        for i in range(int(fps * duration)):
            # Create black frame
            frame = np.zeros((height, width, 3), dtype=np.uint8)

            # Draw moving rectangles
            cv2.rectangle(frame, (i*10 % (width-100), 100), (i*10 % (width-100) + 100, 200), (0, 0, 255), -1)
            cv2.rectangle(frame, (width - (i*8 % (width-80)), 300), (width - (i*8 % (width-80)) + 80, 380), (0, 255, 0), -1)

            # Write frame
            out.write(frame)

        out.release()
        print(f"Created test video: {test_video_path}")
        video_path = test_video_path

    print(f"\nUsing video: {video_path}")

    # Run all AI features
    print("\nRunning Object Detection...")
    run_object_detection(video_path)

    print("\nRunning Object Tracking...")
    run_object_tracking(video_path)

    print("\nRunning Background Blur...")
    run_background_removal(video_path)

    # Generate a simple background image for background replacement
    background_path = os.path.join("test_media", "background.jpg")
    if not os.path.exists(background_path):
        print(f"Creating a simple background image: {background_path}")
        bg_img = np.ones((480, 640, 3), dtype=np.uint8) * 255  # White background

        # Add a simple gradient (vectorized)
        y_coords = np.arange(480)
        gradient = np.stack([
            np.zeros(480, dtype=np.uint8),
            (y_coords / 480 * 255).astype(np.uint8),
            ((1 - y_coords / 480) * 255).astype(np.uint8)
        ], axis=1)
        bg_img = np.tile(gradient[:, np.newaxis, :], (1, 640, 1))

        cv2.imwrite(background_path, cv2.cvtColor(bg_img, cv2.COLOR_RGB2BGR))

    print("\nRunning Background Replacement...")
    run_background_removal(video_path, background_path)

    print("\nAll AI features have been demonstrated.")
    print("Check the output files in the same directory as the input video.")

if __name__ == "__main__":
    main()
