import cv2 as cv
import dlib
import numpy as np
import os

def match_feature_find_object(query_img, train_img, min_matches, results_file_name):
    # Create an ORB object
    orb = cv.ORB_create(nfeatures=100000)

    features1, des1 = orb.detectAndCompute(query_img, None)
    features2, des2 = orb.detectAndCompute(train_img, None)

    # Create Brute-Force matcher object
    bf = cv.BFMatcher(cv.NORM_HAMMING)
    matches = bf.knnMatch(des1, des2, k=2)

    # Nearest neighbour ratio test to find good matches
    good = []
    good_without_lists = []
    matches = [match for match in matches if len(match) == 2]
    for m, n in matches:
        if m.distance < 0.8 * n.distance:
            good.append([m])
            good_without_lists.append(m)

    if len(good) >= min_matches:
        print(f"good: {len(good)}, {results_file_name}")
        # Draw a polygon around the recognized object
        src_pts = np.float32([features1[m.queryIdx].pt for m in good_without_lists]).reshape(-1, 1, 2)
        dst_pts = np.float32([features2[m.trainIdx].pt for m in good_without_lists]).reshape(-1, 1, 2)

        # Get the transformation matrix
        M, _ = cv.findHomography(src_pts, dst_pts, cv.RANSAC, 5.0)

        # Find the perspective transformation to get the corresponding points
        h, w = query_img.shape[:2]
        pts = np.float32([[0, 0], [0, h - 1], [w - 1, h - 1], [w - 1, 0]]).reshape(-1, 1, 2)
        dst = cv.perspectiveTransform(pts, M)

        train_img = cv.polylines(train_img, [np.int32(dst)], True, (0, 255, 0), 2, cv.LINE_AA)

        # Return the bounding box of the detected object
        x, y, w, h = cv.boundingRect(np.int32(dst))
        return (x, y, w, h), train_img  # Return bounding box and modified image

    else:
        print("Not enough good matches are found - {}/{}".format(len(good), min_matches))
        return None, train_img  # No detection, return None


def main():
    input_folder = "extracted_frames/"
    output_folder = "results_extracted_frames/"
    new_object_size = (300, 300)

    os.makedirs(input_folder, exist_ok=True)
    os.makedirs(output_folder, exist_ok=True)

    orb = cv.ORB_create()
    query_img = cv.imread("img/drone.png")
    query_img = cv.resize(query_img, new_object_size)
    features = orb.detect(query_img, None)
    f_img = cv.drawKeypoints(query_img, features, None, color=(0, 255, 0), flags=0)
    # cv.imwrite(os.path.join(output_folder, "drone.png"), f_img)  # Save the image with keypoints

    # Initialize dlib correlation tracker
    tracker = dlib.correlation_tracker()
    is_tracking = False  # Flag to indicate if we are tracking

    # Sort the filenames to process them in sequence
    filenames = sorted([f for f in os.listdir(input_folder) if f.endswith("png")])

    # Previous bounding box (if the object is detected)
    prev_bbox = None

    for filename in filenames:
        img = cv.imread(os.path.join(input_folder, filename))

        if not is_tracking:  # Try to detect using feature matching
            bbox, modified_img = match_feature_find_object(query_img, img, 100, os.path.join(output_folder, filename))

            if bbox is not None:
                # Object detected, initialize the dlib tracker
                x, y, w, h = bbox
                tracker.start_track(img, dlib.rectangle(x, y, x + w, y + h))
                is_tracking = True
                prev_bbox = bbox
            else:
                # No detection, use the image as is
                modified_img = img

        else:  # If we are tracking
            # Update the tracker and get the new position
            tracker.update(img)
            pos = tracker.get_position()
            x1, y1, x2, y2 = int(pos.left()), int(pos.top()), int(pos.right()), int(pos.bottom())
            print(x1, y1, x2, y2)

            # Draw the tracking bounding box
            cv.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)

            # If the bounding box is valid, keep tracking
            if (x2 - x1) > 0 and (y2 - y1) > 0:
                prev_bbox = (x1, y1, x2 - x1, y2 - y1)
                modified_img = img
            else:
                # If tracking fails, stop tracking and try detection again in the next frame
                is_tracking = False
                modified_img = img

        # Save the result as a PNG image
        # cv.imwrite(os.path.join(output_folder, filename), modified_img)


if __name__ == "__main__":
    main()