import cv2
import numpy as np
import joblib
from skimage.feature import hog
import imutils
from imutils.object_detection import non_max_suppression
from skimage.transform import pyramid_gaussian
import asyncio
import concurrent.futures

def sliding_window(image, stepSize, windowSize):
    # Standard (synchronous) sliding window generator (not used in async version)
    for y in range(0, image.shape[0] - windowSize[1] + 1, stepSize):
        for x in range(0, image.shape[1] - windowSize[0] + 1, stepSize):
            yield (x, y, image[y:y + windowSize[1], x:x + windowSize[0]])

# Helper function: generate sliding windows for a given y-range
def generate_windows_for_y_range(image, stepSize, windowSize, y_start, y_end):
    windows = []
    # Ensure y_end does not exceed allowed range
    for y in range(y_start, min(y_end, image.shape[0] - windowSize[1] + 1), stepSize):
        for x in range(0, image.shape[1] - windowSize[0] + 1, stepSize):
            window = image[y:y + windowSize[1], x:x + windowSize[0]]
            # Only add if the window has the desired size
            if window.shape[0] == windowSize[1] and window.shape[1] == windowSize[0]:
                windows.append((x, y, window))
    return windows

# Asynchronous sliding window generator using 5 workers
async def async_sliding_window(image, stepSize, windowSize, num_workers=5):
    loop = asyncio.get_running_loop()
    height = image.shape[0] - windowSize[1] + 1
    segment_size = max(1, height // num_workers)
    tasks = []
    with concurrent.futures.ThreadPoolExecutor(max_workers=num_workers) as executor:
        for i in range(0, height, segment_size):
            y_start = i
            y_end = i + segment_size
            tasks.append(loop.run_in_executor(
                executor,
                generate_windows_for_y_range,
                image, stepSize, windowSize, y_start, y_end
            ))
        results = await asyncio.gather(*tasks)
    # Flatten the list of lists into a single list of windows
    windows = [win for sublist in results for win in sublist]
    return windows

# HOG parameters (must match training parameters)
hog_params = {
    'orientations': 9,
    'pixels_per_cell': (8, 8),
    'cells_per_block': (2, 2),
    'block_norm': 'L2-Hys'
}

# Detection parameters
windowSize = (64, 64)
hogSize = (128, 128)  # Size to which each window is resized for HOG extraction.
stepSize = 8          # Sliding window step size.
scale_factor = 1.1    # Pyramid downscale factor.
detection_threshold = 0.9  # Minimum probability to consider a detection valid.

# Load the pre-trained model
model = joblib.load('svm_natural_images_model.pkl')
if hasattr(model, 'best_estimator_'):
    model = model.best_estimator_

# Initializer for each worker process
def init_worker(_model, _hog_params, _hogSize, _detection_threshold, _windowSize):
    global model, hog_params, hogSize, detection_threshold, windowSize
    model = _model
    hog_params = _hog_params
    hogSize = _hogSize
    detection_threshold = _detection_threshold
    windowSize = _windowSize

# Process a single window: extract features and run prediction
def process_window(args):
    window, x, y, scale_ratio = args
    if window.dtype != np.uint8:
        window = (window * 255).astype(np.uint8)
    # Convert to grayscale and resize for HOG
    gray = cv2.cvtColor(window, cv2.COLOR_BGR2GRAY)
    gray = cv2.resize(gray, hogSize)
    features = hog(gray, **hog_params)
    features = features.reshape(1, -1)
    pred = model.predict(features)
    prob = model.predict_proba(features)[0]
    if 'person' in model.classes_:
        person_index = list(model.classes_).index('person')
        person_prob = prob[person_index]
    else:
        person_prob = 0
    if pred[0] == 'person' and person_prob > detection_threshold:
        startX = int(x * scale_ratio)
        startY = int(y * scale_ratio)
        endX = int((x + windowSize[0]) * scale_ratio)
        endY = int((y + windowSize[1]) * scale_ratio)
        return (startX, startY, endX, endY)
    return None

# Main asynchronous detection function
async def async_detection(image):
    detections = []
    loop = asyncio.get_running_loop()
    tasks = []
    # ProcessPoolExecutor for heavy computation (up to 10 processes)
    with concurrent.futures.ProcessPoolExecutor(
            max_workers=10,
            initializer=init_worker,
            initargs=(model, hog_params, hogSize, detection_threshold, windowSize)
        ) as process_executor:
        # Iterate over each level in the image pyramid
        for resized in pyramid_gaussian(image, downscale=scale_factor, channel_axis=-1):
            scale_ratio = image.shape[1] / float(resized.shape[1])
            # Asynchronously generate sliding