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):
    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]])


# HOG parameters (should match those used during training)
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)  # Must match the window size used for training HOG features.
stepSize = 8  # Step size for sliding window.
scale_factor = 1.1  # Scaling factor for image pyramid.
detection_threshold = 0.9  # Minimum probability to consider a detection valid.

# Load the pre-trained model (grid search object containing the best estimator)
model = joblib.load('svm_natural_images_model.pkl')
# If the saved object is a GridSearchCV, extract the best estimator.
if hasattr(model, 'best_estimator_'):
    model = model.best_estimator_


# This initializer will be run in each worker process to set globals.
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


def process_window(args):
    # Unpack arguments
    window, x, y, scale_ratio = args
    # Ensure proper datatype
    if window.dtype != np.uint8:
        window = (window * 255).astype(np.uint8)
    # Convert to grayscale and resize for HOG extraction
    gray = cv2.cvtColor(window, cv2.COLOR_BGR2GRAY)
    gray = cv2.resize(gray, hogSize)
    # Compute HOG features
    features = hog(gray, **hog_params)
    features = features.reshape(1, -1)
    # Make predictions
    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
    # Return bounding box if detection meets the criteria
    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


async def async_detection(image):
    detections = []
    loop = asyncio.get_running_loop()
    tasks = []
    # Set up the ProcessPoolExecutor with max 10 processes and initializer
    with concurrent.futures.ProcessPoolExecutor(
            max_workers=20,
            initializer=init_worker,
            initargs=(model, hog_params, hogSize, detection_threshold, windowSize)
    ) as executor:
        for resized in pyramid_gaussian(image, downscale=scale_factor, channel_axis=-1):
            scale_ratio = image.shape[1] / float(resized.shape[1])
            for (x, y, window) in sliding_window(resized, stepSize, windowSize):
                if window.shape[0] != windowSize[1] or window.shape[1] != windowSize[0]:
                    continue
                args = (window, x, y, scale_ratio)
                tasks.append(loop.run_in_executor(executor, process_window, args))
        results = await asyncio.gather(*tasks)
    for result in results:
        if result is not None:
            detections.append(result)
    return detections


async def main():
    # Load the test image (update the filename as needed)
    image = cv2.imread("test.jpg")
    if image is None:
        print("Failed to load image!")
        return

    orig = image.copy()
    # Run the asynchronous detection
    detections = await async_detection(image)

    # Apply non-maxima suppression to reduce overlapping boxes
    if len(detections) > 0:
        detections = np.array(detections)
        pick = non_max_suppression(detections, probs=None, overlapThresh=0.3)
        # Draw the final bounding boxes on the original image.
        for (xA, yA, xB, yB) in pick:
            cv2.rectangle(orig, (xA, yA), (xB, yB), (0, 255, 0), 2)

    # Display the detections
    cv2.imshow("Detections", orig)
    cv2.waitKey(0)
    cv2.destroyAllWindows()


if __name__ == '__main__':
    asyncio.run(main())