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

# === Helper: Rotated rectangle corners ===
def get_rotated_rect_corners(x, y, w, h, rotation_deg):
    rot_rad = np.deg2rad(rotation_deg)
    cos_r = np.cos(rot_rad)
    sin_r = np.sin(rot_rad)

    cx = x + w/2
    cy = y + h/2

    local_corners = np.array([
        [-w/2, -h/2],
        [ w/2, -h/2],
        [ w/2,  h/2],
        [-w/2,  h/2]
    ])

    R = np.array([[cos_r, -sin_r],
                  [sin_r,  cos_r]])

    rotated_corners = np.dot(local_corners, R.T)
    corners = rotated_corners + np.array([cx, cy])
    return corners.astype(np.float32)

# === Preprocessing ===
def preprocess_gray_clahe(img):
    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
    return clahe.apply(gray)

# === Feature detectors ===
def get_detector(detector_name):
    if detector_name == "SIFT":
        return cv2.SIFT_create(nfeatures=5000)
    elif detector_name == "ORB":
        return cv2.ORB_create(5000)
    elif detector_name == "BRISK":
        return cv2.BRISK_create()
    elif detector_name == "AKAZE":
        return cv2.AKAZE_create()
    elif detector_name == "KAZE":
        return cv2.KAZE_create()
    else:
        return None

def detect_and_match(img1_gray, img2_gray, detector_name, ratio_thresh=0.78):
    detector = get_detector(detector_name)
    kp1, des1 = detector.detectAndCompute(img1_gray, None)
    kp2, des2 = detector.detectAndCompute(img2_gray, None)

    if detector_name in ["SIFT", "KAZE"]:
        matcher = cv2.BFMatcher(cv2.NORM_L2)
    else:
        matcher = cv2.BFMatcher(cv2.NORM_HAMMING)

    matches = matcher.knnMatch(des1, des2, k=2)
    good = []
    for m, n in matches:
        if m.distance < ratio_thresh * n.distance:
            good.append(m)
    return kp1, kp2, good

# === Main processing ===
def process_images(flat_img, persp_img, json_file):
    if flat_img is None or persp_img is None or json_file is None:
        return [None]*6

    # Load JSON
    try:
        data = json.load(open(json_file.name))
    except Exception as e:
        print("JSON read error:", e)
        return [None]*6

    roi = data["printAreas"][0]
    roi_x = roi["position"]["x"]
    roi_y = roi["position"]["y"]
    roi_w = roi["width"]
    roi_h = roi["height"]
    roi_rot_deg = roi["rotation"]

    # Preprocess images
    flat_gray = preprocess_gray_clahe(flat_img)
    persp_gray = preprocess_gray_clahe(persp_img)

    detectors = ["SIFT", "ORB", "BRISK", "AKAZE", "KAZE"]
    results = []

    for det in detectors:
        kp1, kp2, good_matches = detect_and_match(flat_gray, persp_gray, det)
        if len(good_matches) < 4:
            print(f"Not enough matches for {det}")
            results.append(None)
            continue

        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)

        # ROI corners
        roi_corners_flat = get_rotated_rect_corners(roi_x, roi_y, roi_w, roi_h, roi_rot_deg)
        roi_corners_persp = cv2.perspectiveTransform(roi_corners_flat.reshape(-1,1,2), H).reshape(-1,2)

        # Draw ROI
        flat_out = flat_img.copy()
        persp_out = persp_img.copy()
        cv2.polylines(flat_out, [roi_corners_flat.astype(int)], True, (255,0,0), 3)
        cv2.polylines(persp_out, [roi_corners_persp.astype(int)], True, (0,255,0), 3)

        results.append([flat_out, persp_out])

    return results  # List of [flat_out, persp_out] for each detector

# === Gradio Interface ===
def wrap_gradio(flat_img, persp_img, json_file):
    outputs = process_images(flat_img, persp_img, json_file)
    # Flatten the outputs for Gallery display
    gallery_images = []
    for item in outputs:
        if item is not None:
            gallery_images.extend([item[0], item[1]])
    return gallery_images

iface = gr.Interface(
    fn=wrap_gradio,
    inputs=[
        gr.Image(type="numpy", label="Flat Image"),
        gr.Image(type="numpy", label="Perspective Image"),
        gr.File(type="filepath", label="JSON File")
    ],
    outputs=[
        gr.Gallery(label="Results (Flat + Perspective per Detector)")
    ],
    title="Feature Detection with ROI Projection",
    description="Shows SIFT, ORB, BRISK, AKAZE, KAZE feature-based ROI projections. Each detector outputs Flat and Perspective images."
)

iface.launch()