"""
https://github.com/kornia/kornia
"""
import cv2
import torch
import kornia as K
import kornia.feature as KF
import matplotlib.pyplot as plt
from kornia_moons.viz import draw_LAF_matches


def match_keyboards(img_path1, img_path2):
    # 1. 加载图片并转换为 Tensor (Kornia 处理的是 Tensor)
    # 键盘细节丰富，建议保持较高分辨率，不要缩放得太厉害
    img1 = K.image_to_tensor(cv2.imread(img_path1), keepdim=False).float() / 255.0
    img2 = K.image_to_tensor(cv2.imread(img_path2), keepdim=False).float() / 255.0

    # 转换为灰度图用于特征提取
    img1_gray = K.color.rgb_to_grayscale(K.color.bgr_to_rgb(img1))
    img2_gray = K.color.rgb_to_grayscale(K.color.bgr_to_rgb(img2))

    # 2. 定义特征提取器 (KeyNet + AffNet + HardNet)
    # 这是一套非常强大的组合，对键盘这种重复纹理有很好的识别力
    num_features = 2000  # 键盘字符多，建议点数设多一点
    matcher = KF.LocalFeatureMatcher(
        KF.KeyNetAffNetHardNet(num_features, pretrained=True),
        KF.DescriptorMatcher('adalam', torch.device('cpu'))  # 'adalam' 是目前非常快的离群点过滤匹配算法
    )

    # 3. 执行匹配
    input_dict = {"image0": img1_gray, "image1": img2_gray}
    with torch.no_grad():
        correspondences = matcher(input_dict)

    # 提取匹配后的坐标点
    mkpts0 = correspondences['keypoints0'].cpu().numpy()
    mkpts1 = correspondences['keypoints1'].cpu().numpy()

    print(f"找到匹配点数量: {len(mkpts0)}")

    # 4. 可视化结果
    # 我们用 OpenCV 把两张图拼在一起看连线
    img1_cv = cv2.cvtColor(cv2.imread(img_path1), cv2.COLOR_BGR2RGB)
    img2_cv = cv2.cvtColor(cv2.imread(img_path2), cv2.COLOR_BGR2RGB)

    # 简单的可视化：画出前 50 个最强的匹配
    fig, ax = plt.subplots(figsize=(15, 8))
    draw_LAF_matches(
        KF.laf_from_center_scale_ori(correspondences['keypoints0'].view(1, -1, 2)),
        KF.laf_from_center_scale_ori(correspondences['keypoints1'].view(1, -1, 2)),
        torch.arange(len(mkpts0)).view(1, -1, 1),
        img1_cv,
        img2_cv,
        draw_dict={'inliers_fallback': True}
    )
    plt.show()

    # 5. 判断逻辑
    # 如果匹配点数量很多（比如 > 100）且连线平行度高，则极有可能是同一个键盘
    if len(mkpts0) > 50:
        print("结论：两张图片匹配度极高，很可能是同一个键盘或极其相似。")
    else:
        print("结论：匹配点过少，可能是不同键盘或拍摄角度偏差过大。")

# 使用示例
match_keyboards(
    'keyboard_left.jpg',
    'keyboard_right.jpg'
)