import cv2
import numpy as np
import gradio as gr


# Define a blurring function.
def blur(face, factor=3):

    h, w = face.shape[:2]

    if factor < 1:
        factor = 1  # Maximum blurring
    if factor > 5:
        factor = 5  # Minimal blurring

    # Kernel size.
    w_k = int(w / factor)
    h_k = int(h / factor)

    # Insure kernel is an odd number.
    if w_k % 2 == 0:
        w_k += 1
    if h_k % 2 == 0:
        h_k += 1

    blurred = cv2.GaussianBlur(face, (int(w_k), int(h_k)), 0, 0)
    return blurred


def pixelate(roi, pixels=16):

    # Size of region to pixelate.
    roi_h, roi_w = roi.shape[:2]

    if roi_h > pixels and roi_w > pixels:
        # Resize input ROI to the (small) pixelated size.
        roi_small = cv2.resize(roi, (pixels, pixels), interpolation=cv2.INTER_LINEAR)

        # Now enlarge the pixelated ROI to fill the size of the original ROI.
        roi_pixelated = cv2.resize(roi_small, (roi_w, roi_h), interpolation=cv2.INTER_NEAREST)
    else:
        roi_pixelated = roi

    return roi_pixelated


def face_blur_ellipse_pixelate(image, net, detect_threshold=0.9, factor=3, pixels=10, write_mask=False):

    img = image.copy()
    img_out = img.copy()
    elliptical_mask = np.zeros(img.shape, dtype=img.dtype)

    # Prepare image and perform inference.
    blob = cv2.dnn.blobFromImage(img, scalefactor=1.0, size=(300, 300), mean=[104, 117, 123])
    net.setInput(blob)
    detections = net.forward()

    h, w = img.shape[:2]
    for i in range(detections.shape[2]):
        confidence = detections[0, 0, i, 2]
        if confidence > detect_threshold:

            # Extract the bounding box coordinates from the detection.
            box = detections[0, 0, i, 3:7] * np.array([w, h, w, h])
            (x1, y1, x2, y2) = box

            # The face is defined by the bounding rectangle from the detection.
            face = img[int(y1) : int(y2), int(x1) : int(x2), :]

            # Blur the rectangular area defined by the bounding box.
            face = blur(face, factor=factor)

            # Pixelate the blurred face.
            face = pixelate(face, pixels=pixels)

            # Copy the blurred/pixelated face to the output image.
            img_out[int(y1) : int(y2), int(x1) : int(x2), :] = face

            # Specify the elliptical parameters directly from the bounding box coordinates.
            e_center = (x1 + (x2 - x1) / 2, y1 + (y2 - y1) / 2)
            e_size = (x2 - x1, y2 - y1)
            e_angle = 0.0

            # Create an elliptical mask.
            elliptical_mask = cv2.ellipse(
                elliptical_mask, (e_center, e_size, e_angle), (255, 255, 255), -1, cv2.LINE_AA
            )
            # Apply the elliptical mask.
            np.putmask(img, elliptical_mask, img_out)

    if write_mask:
        cv2.imwrite("elliptical_mask.jpg", elliptical_mask)

    return img


# img1_epb = face_blur_ellipse_pixelate(img1, net, factor=3.5, pixels=15)
# img2_epb = face_blur_ellipse_pixelate(img2, net, factor=2, pixels=10)


# cv2.imshow('image-out', img1_epb)
# cv2.waitKey(0)
# cv2.imshow('image-out', img2_epb)
# cv2.waitKey(0)
# cv2.destroyAllWindows()


def process_video(input_video_path, detect_threshold=0.9, blur_factor=3, pixel_size=10):
    
    # Load the DNN model.
    modelFile = "./model/res10_300x300_ssd_iter_140000.caffemodel"
    configFile = "./model/deploy.prototxt"

    # Read the model and create a network object.
    net = cv2.dnn.readNetFromCaffe(prototxt=configFile, caffeModel=modelFile)
    # Initialize video capture.
    cap = cv2.VideoCapture(input_video_path)
    if not cap.isOpened():
        print(f"Error: Cannot open video file {input_video_path}")
        return

    # 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))
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")  # You can change the codec as needed.
    output_video_path = "output_video.mp4"

    # Initialize video writer.
    out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))

    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    current_frame = 0

    print(f"Processing video... Total frames: {frame_count}")

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

        # Process the frame.
        processed_frame = face_blur_ellipse_pixelate(frame, net, detect_threshold, blur_factor, pixel_size)

        # cv2.imshow('output', processed_frame)

        # Write the processed frame to the output video.
        out.write(processed_frame)

        current_frame += 1
        if current_frame % 30 == 0 or current_frame == frame_count:
            print(f"Processed {current_frame}/{frame_count} frames")

    # Release resources.
    cap.release()
    out.release()
    print(f"Processing complete. Output saved to {output_video_path}")
    return output_video_path


# input_video = "./input-video.mp4"

# process_video(input_video)


# Define the Gradio interface
iface = gr.Interface(
    fn=process_video,
    inputs=[
        gr.Video(label="Input Video", autoplay=True),
        gr.Slider(minimum=0.5, maximum=0.99, value=0.9, label="Detection Confidence Threshold"),
        gr.Slider(minimum=1, maximum=5, step=0.5, value=3, label="Blur Factor"),
        gr.Slider(minimum=5, maximum=50, step=1, value=10, label="Pixelation Size"),
    ],
    outputs=[gr.Video(label="Processed Video", autoplay=True)],
    title="Face Blurring with OpenCV",
    description="Upload a video file (MP4 or AVI). The app will detect faces, blur and pixelate them, and provide a processed video for download.",
    examples=[["./input-video.mp4", 0.9, 3, 10], ["./man-woman.mp4", 0.85, 4, 12]],
    cache_examples=False,
)

if __name__ == "__main__":
    iface.launch()