import os
import cv2 as cv
from detector import Detector
import matplotlib.pyplot as plt

SHOW_FAILED = False

def show_image(img, title, colorspace, results_file_name, show=False):
    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 main():
    input_folder = "data/video-frames/"
    success_output_folder = "data/success-video-frames/"
    failures_output_folder = "data/failures-video-frames/"

    os.makedirs(input_folder, exist_ok=True)
    os.makedirs(success_output_folder, exist_ok=True)
    os.makedirs(failures_output_folder, exist_ok=True)

    target_image = cv.imread("cropped_image.png")

    detector = Detector(target_image, min_matches=4, angle_tolerance=20, length_ratio_tolerance=2.5)

    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, modified_img = detector.match_feature_find_object(img)
        if bbox is not None:
            show_image(modified_img, "Success feature matching", "RGB", os.path.join(success_output_folder, filename), show=True)
        else:
            show_image(modified_img, "Failed feature matching", "RGB", os.path.join(failures_output_folder, filename), show=SHOW_FAILED)

        

if __name__ == "__main__":
    main()