import cv2
import numpy as np
import json
import gradio as gr
import os

# ---------------- Helper functions ----------------
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)

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

    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]
    ])

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

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

def detect_and_match(img1_gray, img2_gray, method="SIFT", ratio_thresh=0.78):
    if method == "SIFT":
        sift = cv2.SIFT_create(nfeatures=5000)
        kp1, des1 = sift.detectAndCompute(img1_gray, None)
        kp2, des2 = sift.detectAndCompute(img2_gray, None)
        matcher = cv2.BFMatcher(cv2.NORM_L2)
    elif method == "ORB":
        orb = cv2.ORB_create(5000)
        kp1, des1 = orb.detectAndCompute(img1_gray, None)
        kp2, des2 = orb.detectAndCompute(img2_gray, None)
        matcher = cv2.BFMatcher(cv2.NORM_HAMMING)
    elif method == "BRISK":
        brisk = cv2.BRISK_create()
        kp1, des1 = brisk.detectAndCompute(img1_gray, None)
        kp2, des2 = brisk.detectAndCompute(img2_gray, None)
        matcher = cv2.BFMatcher(cv2.NORM_HAMMING)
    elif method == "KAZE":
        kaze = cv2.KAZE_create()
        kp1, des1 = kaze.detectAndCompute(img1_gray, None)
        kp2, des2 = kaze.detectAndCompute(img2_gray, None)
        matcher = cv2.BFMatcher(cv2.NORM_L2)
    elif method == "AKAZE":
        akaze = cv2.AKAZE_create()
        kp1, des1 = akaze.detectAndCompute(img1_gray, None)
        kp2, des2 = akaze.detectAndCompute(img2_gray, None)
        matcher = cv2.BFMatcher(cv2.NORM_HAMMING)
    else:
        return None, None, []

    if des1 is None or des2 is None:
        return None, None, []

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

# ---------------- Main Homography Function ----------------
def homography_all_detectors(flat_file, persp_file, json_file):
    flat_img = cv2.imread(flat_file)
    persp_img = cv2.imread(persp_file)
    mockup = json.load(open(json_file.name))

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

    flat_gray = preprocess_gray_clahe(flat_img)
    persp_gray = preprocess_gray_clahe(persp_img)

    methods = ["SIFT", "ORB", "BRISK", "KAZE", "AKAZE"]
    gallery_images = []
    download_files = []

    for method in methods:
        kp1, kp2, good_matches = detect_and_match(flat_gray, persp_gray, method=method)

        if kp1 is None or kp2 is None or len(good_matches) < 4:
            continue  # skip if no matches

        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)

        if H is None:
            continue

        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)

        persp_debug = persp_img.copy()
        cv2.polylines(persp_debug, [roi_corners_persp.astype(int)], True, (0,255,0), 2)
        for (px, py) in roi_corners_persp:
            cv2.circle(persp_debug, (int(px), int(py)), 5, (255,0,0), -1)

        # Convert BGR -> RGB for display
        result_rgb = cv2.cvtColor(persp_debug, cv2.COLOR_BGR2RGB)

        # Save result for download
        file_name = f"result_{method.lower()}.png"
        cv2.imwrite(file_name, cv2.cvtColor(result_rgb, cv2.COLOR_RGB2BGR))

        gallery_images.append((result_rgb, f"{method} Result"))
        download_files.append(file_name)

    # return gallery + 5 download files (pad with None if less)
    while len(download_files) < 5:
        download_files.append(None)

    return [gallery_images] + download_files[:5]

# ---------------- Gradio UI ----------------
iface = gr.Interface(
    fn=homography_all_detectors,
    inputs=[
        gr.Image(label="Upload Flat Image", type="filepath"),
        gr.Image(label="Upload Perspective Image", type="filepath"),
        gr.File(label="Upload mockup.json", file_types=[".json"])
    ],
    outputs=[
        gr.Gallery(label="Results (per Detector)", show_label=True),
        gr.File(label="Download SIFT Result"),
        gr.File(label="Download ORB Result"),
        gr.File(label="Download BRISK Result"),
        gr.File(label="Download KAZE Result"),
        gr.File(label="Download AKAZE Result")
    ],
    title="Homography ROI Projection with Multiple Feature Detectors",
    description="Upload flat & perspective images with mockup.json. The system will project ROI using SIFT, ORB, BRISK, KAZE, and AKAZE. Each result can be viewed and downloaded."
)

iface.launch()