import os
import cv2 as cv
import numpy as np
from detector import Detector
from tracker import Tracker
import matplotlib.pyplot as plt


def show_image(img, title, colorspace, results_file_name, show=False):
    """Displays or saves the processed image."""
    if not show:
        return
    dpi = 96
    figsize = (img.shape[1] / dpi, img.shape[0] / dpi)
    fig, ax = plt.subplots(figsize=figsize, dpi=dpi)
    if colorspace == "RGB":
        plt.imshow(cv.cvtColor(img, cv.COLOR_BGR2RGB), interpolation="spline16")
    if colorspace == "gray":
        plt.imshow(img, cmap="gray")
    plt.title(title, fontsize=12)
    ax.axis("off")
    plt.savefig(results_file_name)
    plt.close()


def draw_bounding_box(img, bbox, color=(0, 255, 0)):
    """Draws a bounding box (as a polygon) on the image."""
    x, y, w, h = bbox
    box_points = np.array([
        [x, y],
        [x + w, y],
        [x + w, y + h],
        [x, y + h]
    ], dtype=np.int32)
    cv.polylines(img, [box_points], isClosed=True, color=color, thickness=2, lineType=cv.LINE_AA)


def main():
    input_folder = "data/video-frames/"
    output_folder = "data/detect_and_track/results/"

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

    query_img = cv.imread("cropped_image.png")
    detector = Detector(query_img, min_matches=4, angle_tolerance=20, length_ratio_tolerance=2.5)
    tracker = Tracker()

    filenames = sorted([f for f in os.listdir(input_folder) if f.endswith("png")])

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

        bbox, detected_img = detector.match_feature_find_object(img)
        if bbox is not None:
            tracker.start_tracking(img, bbox)
            draw_bounding_box(img, bbox, color=(0, 255, 0))
            modified_img = detected_img
        else:
            if tracker.is_tracking:
                bbox = tracker.update(img)
                if bbox:
                    draw_bounding_box(img, bbox, color=(255, 0, 0))
                else:
                    tracker.is_tracking = False
            modified_img = img

        show_image(modified_img, "Feature matching and object recognition", "RGB", os.path.join(output_folder, filename), show=(bbox is not None))


if __name__ == "__main__":
    main()