import argparse
import numpy as np
import cv2
import axengine as axe

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", type=str, required=True, help="Path to the axmodel")
    parser.add_argument("--img1", type=str, required=True, help="Path to the first image")
    parser.add_argument("--img2", type=str, required=True, help="Path to the second image")
    parser.add_argument("--output", type=str, default="matches.jpg", help="The output image directory")
    parser.add_argument("--threshold", type=float, default=0.005, help="The keypoint threshold")
    parser.add_argument("--max_points", type=int, default=100, help="The max num for keypoints")
    return parser.parse_args()

def preprocess_image(path: str, h: int, w: int):
    img = cv2.imread(path)
    raw_h, raw_w = img.shape[:2]

    if (raw_h, raw_w) != (h, w):
        img = cv2.resize(img, (w, h))
        scale_h = raw_h / h
        scale_w = raw_w / w
    else:
        scale_h = 1.0
        scale_w = 1.0

    img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    img_tensor = img_gray.astype(np.float32) / 255.0
    img_tensor = img_tensor[None, None, :, :]  # -> (1, 1, H, W)

    return img_tensor, img, (scale_h, scale_w)

def get_keypoints(score_map, threshold):
    row, col = np.where(score_map > threshold)  # y, x
    if len(row) == 0:
        return np.zeros((0, 2), dtype=np.float32), np.zeros((0,), dtype=np.float32)
    
    scores = score_map[row, col]
    keypoints = np.stack([col, row], axis=1).astype(np.float32)
    return keypoints, scores

def get_descriptors(kp, desc_map):
    if len(kp) == 0:
        return np.zeros((0, 256), dtype=np.float32)
    
    c, h, w = desc_map.shape
    x = kp[:, 0] / 8.0
    y = kp[:, 1] / 8.0

    x0 = np.floor(x).astype(np.int32)
    x1 = x0 + 1
    y0 = np.floor(y).astype(np.int32)
    y1 = y0 + 1

    x0 = np.clip(x0, 0, w - 1)
    x1 = np.clip(x1, 0, w - 1)
    y0 = np.clip(y0, 0, h - 1)
    y1 = np.clip(y1, 0, h - 1)

    wa = (x1 - x) * (y1 - y)
    wb = (x1 - x) * (y - y0)
    wc = (x - x0) * (y1 - y)
    wd = (x - x0) * (y - y0)

    wa = wa[None, :]
    wb = wb[None, :]
    wc = wc[None, :]
    wd = wd[None, :]

    Q_tl = desc_map[:, y0, x0]
    Q_bl = desc_map[:, y1, x0]
    Q_tr = desc_map[:, y0, x1]
    Q_br = desc_map[:, y1, x1]

    sampled = (Q_tl * wa + Q_bl * wb + Q_tr * wc + Q_br * wd)
    descriptors = sampled.T
    norm = np.linalg.norm(descriptors, axis=1, keepdims=True)
    descriptors = descriptors / (norm + 1e-6)
    
    return descriptors.astype(np.float32)

def infer(model: str, img1_path: str, img2_path: str, output: str, threshold: float, max_points: int):
    session = axe.InferenceSession(model)

    # superpoint only have one input
    input_name = session.get_inputs()[0].name  # get model input node name
    input_shape = session.get_inputs()[0].shape  # get model input shape (1, 1, H, W)
    target_h, target_w = input_shape[2], input_shape[3]
    print(f"Inference resolution: {target_w}x{target_h}")

    # preprocess images
    input_tensor1, img1, scale1 = preprocess_image(img1_path, target_h, target_w)
    input_tensor2, img2, scale2 = preprocess_image(img2_path, target_h, target_w)

    res1 = session.run(None, {input_name: input_tensor1})
    res2 = session.run(None, {input_name: input_tensor2})

    # [1,480,640], [1,256,60,80]
    score_map1, desc1_map = res1[0], res1[1]
    score_map2, desc2_map = res2[0], res2[1]

    keypoints1, scores1 = get_keypoints(score_map1[0], threshold)
    keypoints2, scores2 = get_keypoints(score_map2[0], threshold)

    print(f"Found {len(keypoints1)} keypoints in image 1")
    print(f"Found {len(keypoints2)} keypoints in image 2")

    if len(keypoints1) > max_points:
        idx = np.argsort(scores1)[::-1][:max_points]
        keypoints1 = keypoints1[idx]
        scores1 = scores1[idx]
    if len(keypoints2) > max_points:
        idx = np.argsort(scores2)[::-1][:max_points]
        keypoints2 = keypoints2[idx]
        scores2 = scores2[idx]

    desc1 = get_descriptors(keypoints1, desc1_map[0])
    desc2 = get_descriptors(keypoints2, desc2_map[0])

    bf = cv2.BFMatcher(cv2.NORM_L2, crossCheck=True)
    matches = bf.match(desc1, desc2)
    matches = sorted(matches, key=lambda x: x.distance)

    points1 = [cv2.KeyPoint(x, y, 1) for x, y in keypoints1]
    points2 = [cv2.KeyPoint(x, y, 1) for x, y in keypoints2]

    match_img = cv2.drawMatches(
        img1, points1,
        img2, points2,
        matches, None,
        flags=cv2.DrawMatchesFlags_NOT_DRAW_SINGLE_POINTS,
        matchColor=(0, 255, 0)
    )

    # if len(matches) > 4:
    #     pts1 = np.float32([keypoints1[m.queryIdx] for m in matches]).reshape(-1, 1, 2)
    #     pts2 = np.float32([keypoints2[m.trainIdx] for m in matches]).reshape(-1, 1, 2)

    #     H, mask = cv2.findHomography(pts1, pts2, cv2.RANSAC, 3.0)

    #     if mask is not None:
    #         matches_mask = mask.ravel().tolist()
    #         inlier_matches = [m for i, m in enumerate(matches) if matches_mask[i]]
    #         print(f"Inliers: {len(inlier_matches)} / {len(matches)}")

    #         inlier_img = cv2.drawMatches(
    #             img1, points1,
    #             img2, points2,
    #             inlier_matches, None,
    #             flags=cv2.DrawMatchesFlags_NOT_DRAW_SINGLE_POINTS,
    #             matchColor=(0, 255, 0)
    #         )
    #         cv2.imwrite("inliers_" + output, inlier_img)

    cv2.imwrite(output, match_img)
    print(f"Result saved to {output}")

def main():
    args = parse_args()
    infer(args.model, args.img1, args.img2, args.output, args.threshold, args.max_points)

if __name__ == '__main__':
    main()