#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2024/5/31 0031 上午 9:55
# @Author  : liudan
# @File    : image_match_demo.py
# @Software: pycharm
#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2024/5/23 0023 下午 2:25
# @Author  : liudan
# @File    : demo.py
# @Software: pycharm

import cv2 as cv
import cv2
import numpy as np
from loguru import logger
import os

from superpoint_superglue_deployment import Matcher
from datetime import datetime
import random
# timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
#import image_similarity_count


def registration_demo():
    query_image = cv2.imread('./data/images/frame_002100_small.jpg') #原图2160*3840
    ref_image = cv2.imread('./data/images/frame_002200_small.jpg')

    new_size = (800, 1420) # W,H
    query_image_resize = cv2.resize(query_image,new_size)
    ref_image_resize = cv2.resize(ref_image,new_size)

    query_gray = cv2.cvtColor(query_image_resize, cv2.COLOR_BGR2GRAY)
    ref_gray = cv2.cvtColor(ref_image_resize, cv2.COLOR_BGR2GRAY)

    superglue_matcher = Matcher(
        {
            "superpoint": {
                "input_shape": (-1, -1),
                "keypoint_threshold": 0.003,
            },
            "superglue": {
                "match_threshold": 0.5,
            },
            "use_gpu": True,
        }
    )
    query_kpts, ref_kpts, _, _, matches = superglue_matcher.match(query_gray, ref_gray)
    M, mask = cv2.findHomography(
        np.float64([query_kpts[m.queryIdx].pt for m in matches]).reshape(-1, 1, 2),
        np.float64([ref_kpts[m.trainIdx].pt for m in matches]).reshape(-1, 1, 2),
        method=cv2.USAC_MAGSAC,
        ransacReprojThreshold=5.0,
        maxIters=10000,
        confidence=0.95,
    )
    logger.info(f"number of inliers: {mask.sum()}")
    matches = np.array(matches)[np.all(mask > 0, axis=1)]
    matches = sorted(matches, key=lambda match: match.distance)
    matched_image = cv2.drawMatches(
        query_image_resize,
        query_kpts,
        ref_image_resize,
        ref_kpts,
        matches[:50],
        None,
        flags=2,
    )

    demo_image_path = './data/similarity_image/demo_result/'
    # match_file_name = f"match_image_{idx1}_{idx2}_{iteration}.jpg"
    # cv2.imwrite(os.path.join(demo_image_path , match_file_name), matched_image)
    cv2.imwrite(os.path.join(demo_image_path + f"matched_image.jpg"), matched_image)

    wrap_image = cv.warpPerspective(query_image_resize, M, (ref_image_resize.shape[1], ref_image_resize.shape[0]))
    wrap_image = cv2.resize(wrap_image,(ref_image.shape[1], ref_image.shape[0]))
    # wrap_filename = f"wrap_image_{idx1}_{idx2}_{iteration}.jpg"
    # cv2.imwrite(os.path.join(demo_image_path, wrap_filename), wrap_image)
    cv2.imwrite(os.path.join(demo_image_path + f"wrap_image.jpg"), wrap_image)

    result_image = cv2.subtract(ref_image, wrap_image)
    # result_file_name = f"result_image_{idx1}_{idx2}_{iteration}.jpg"
    cv2.imwrite(os.path.join(demo_image_path + f"result_image.jpg"), result_image)
    # cv2.imwrite(os.path.join(demo_image_path, result_file_name), result_image)

    return wrap_image



if __name__ == "__main__":
    registration_demo()

    #image_similarity_count.compare_boxes_similarity(image1_path, image2_path, json_file_path, similarity_threshold=0.5)