import gradio as gr
import cv2
import numpy as np
import matplotlib.pyplot as plt
import json
import math
import os

def ransac(image1, image2, detector_type):
    """
    Finds the homography matrix using the RANSAC algorithm with the selected feature detector.
    """
    gray1 = cv2.cvtColor(image1, cv2.COLOR_RGB2GRAY)
    gray2 = cv2.cvtColor(image2, cv2.COLOR_RGB2GRAY)

    if detector_type == "SIFT":
        detector = cv2.SIFT_create()
        matcher = cv2.FlannBasedMatcher(dict(algorithm=1, trees=5), dict(checks=50))
    elif detector_type == "ORB":
        detector = cv2.ORB_create()
        matcher = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)
    elif detector_type == "BRISK":
        detector = cv2.BRISK_create()
        matcher = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)
    elif detector_type == "AKAZE":
        detector = cv2.AKAZE_create()
        matcher = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)
    elif detector_type == "KAZE":
        detector = cv2.KAZE_create()
        matcher = cv2.BFMatcher(cv2.NORM_L2, crossCheck=True)
    else:
        return None

    kp1, des1 = detector.detectAndCompute(gray1, None)
    kp2, des2 = detector.detectAndCompute(gray2, None)

    if des1 is None or des2 is None or len(kp1) < 2 or len(kp2) < 2:
        return None

    try:
        if detector_type == "SIFT":
            matches = matcher.knnMatch(des1, des2, k=2)
            good_matches = []
            if matches:
                for m, n in matches:
                    if m.distance < 0.75 * n.distance:
                        good_matches.append(m)
        else:
            matches = matcher.match(des1, des2)
            good_matches = sorted(matches, key=lambda x: x.distance)
    except cv2.error as e:
        print(f"Error during matching: {e}")
        return None

    if len(good_matches) > 10:
        src_pts = np.float32([kp1[m.queryIdx].pt for m in good_matches]).reshape(-1, 1, 2)
        dst_pts = np.float32([kp2[m.trainIdx].pt for m in good_matches]).reshape(-1, 1, 2)
        H, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 5.0)
        return H
    else:
        return None

def get_bounding_box_points(json_data):
    """
    Extracts and calculates the four corner points of the bounding box, assuming x,y are top-left.
    """
    print_area = json_data['printAreas'][0]
    x = print_area['position']['x']
    y = print_area['position']['y']
    w = print_area['width']
    h = print_area['height']
    rotation_deg = print_area['rotation']

    points = np.float32([
        [0, 0],
        [w, 0],
        [w, h],
        [0, h]
    ]).reshape(-1, 1, 2)

    rotation_rad = math.radians(rotation_deg)
    cos_theta = math.cos(rotation_rad)
    sin_theta = math.sin(rotation_rad)
    rotation_matrix = np.array([
        [cos_theta, -sin_theta],
        [sin_theta,  cos_theta]
    ])

    rotated_points = np.dot(points.reshape(-1, 2), rotation_matrix.T)
    final_points = rotated_points + np.array([x, y])
    return final_points.reshape(-1, 1, 2)

def process_and_plot_all_detectors(image1_np, image2_np, json_file):
    """
    Processes the images with all available detectors and returns image data for display and download.
    """
    if image1_np is None or image2_np is None:
        return [None] * 6

    try:
        with open(json_file.name, 'r') as f:
            data = json.load(f)
    except Exception as e:
        print(f"Error: Could not read JSON file. {e}")
        return [None] * 6

    detectors = ["SIFT", "ORB", "BRISK", "AKAZE", "KAZE"]
    gallery_images = []
    download_files = [None] * 5

    for i, detector_type in enumerate(detectors):
        H = ransac(image1_np, image2_np, detector_type)

        if H is not None:
            box_points = get_bounding_box_points(data)
            
            output_flat_img = image1_np.copy()
            cv2.polylines(output_flat_img, [np.int32(box_points)], isClosed=True, color=(0, 0, 255), thickness=5)

            transformed_box_points = cv2.perspectiveTransform(box_points, H)

            output_perspective_img = image2_np.copy()
            cv2.polylines(output_perspective_img, [np.int32(transformed_box_points)], isClosed=True, color=(0, 0, 255), thickness=5)

            fig, axes = plt.subplots(1, 3, figsize=(18, 6))
            axes[0].imshow(cv2.cvtColor(output_flat_img, cv2.COLOR_BGR2RGB))
            axes[0].set_title(f'Original (Flat) - {detector_type}')
            axes[0].axis('off')

            axes[1].imshow(cv2.cvtColor(image2_np, cv2.COLOR_BGR2RGB))
            axes[1].set_title('Original (Perspective)')
            axes[1].axis('off')

            axes[2].imshow(cv2.cvtColor(output_perspective_img, cv2.COLOR_BGR2RGB))
            axes[2].set_title('Projected Bounding Box')
            axes[2].axis('off')
            
            plt.tight_layout()
            
            file_name = f"result_{detector_type.lower()}.png"
            plt.savefig(file_name)
            plt.close(fig)

            gallery_images.append(file_name)
            download_files[i] = file_name
        else:
            print(f"Warning: Homography matrix could not be found with {detector_type} detector. Skipping this result.")
            # We don't append None to the gallery_images list to avoid the error.
            # download_files[i] remains None, which is handled correctly by gr.File.
    
    return [gallery_images] + download_files

iface = gr.Interface(
    fn=process_and_plot_all_detectors,
    inputs=[
        gr.Image(type="numpy", label="Image 1 (Flat)"),
        gr.Image(type="numpy", label="Image 2 (Perspective)"),
        gr.File(type="filepath", label="JSON File")
    ],
    outputs=[
        gr.Gallery(label="Results"),
        gr.File(label="Download SIFT Result"),
        gr.File(label="Download ORB Result"),
        gr.File(label="Download BRISK Result"),
        gr.File(label="Download AKAZE Result"),
        gr.File(label="Download KAZE Result")
    ],
    title="Homography and Bounding Box Projection with All Detectors",
    description="Upload two images and a JSON file to see the bounding box projection for all 5 feature extraction methods. Each result can be downloaded separately."
)

iface.launch()