#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2024/6/6 0006 上午 9:21
# @Author  : liudan
# @File    : demo_env.py
# @Software: pycharm

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

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

def process_image_with_mask(ref_image, scale_factor, json_ref_path, ref_gray):

    with open(json_ref_path, 'r') as f:
        data = json.load(f)

    shapes = data['shapes']
    shape = shapes[0]
    if shape['shape_type'] == 'polygon':
        coords = [(int(round(x * scale_factor)), int(round(y * scale_factor))) for x, y in shape['points']]
    else:
        coords = []

    mask = np.zeros(ref_gray.shape, dtype=np.uint8) * 255
    pts = np.array(coords, np.int32)
    cv2.fillPoly(mask, [pts], 1)
    ref_gray_masked = cv2.bitwise_and(ref_gray, ref_gray, mask=mask)
    cv2.imwrite('1111111.jpg',ref_gray_masked)
    return ref_gray_masked


def registration_demo(image_dir,demo_image_path, json_ref_path, ref_image_path):
    wraped_image_list =[]
    for filename in os.listdir(image_dir):
        # 检查文件是否是图片（这里假设是.jpg格式）
        try:
            if filename.endswith('.jpg') or filename.endswith('.JPG'):
            # 构建图片的完整路径
                image1_path = os.path.join(image_dir, filename)
                query_image = cv2.imread(image1_path)
                ref_image = cv2.imread(ref_image_path)

                # height, width = ref_image.shape[:2]
                # scale_factor = 0.37
                # query_image_resize = cv2.resize(query_image, dsize=None, fx=scale_factor, fy=scale_factor)
                # ref_image_resize = cv2.resize(ref_image, dsize=None, fx=scale_factor, fy=scale_factor)

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

                query_gray = cv2.cvtColor(query_image, cv2.COLOR_BGR2GRAY)
                ref_gray = cv2.cvtColor(ref_image, cv2.COLOR_BGR2GRAY)

                if os.path.exists(json_ref_path):
                    ref_gray= process_image_with_mask(ref_image, scale_factor, json_ref_path,ref_gray)

                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)
                print(f"{filename}\n")
                matched_query_kpts =[query_kpts[m.queryIdx].pt for m in matches]
                matched_ref_kpts =[ref_kpts[m.trainIdx].pt for m in matches]
                diff_query_ref=np.array(matched_ref_kpts) -np.array(matched_query_kpts)
                # for index in range(len(matched_query_kpts)):
                #     diff_query_ref.append(matched_query_kpts[index]-matched_ref_kpts[index])

                print(matched_query_kpts)
                print(matched_ref_kpts)
                print(diff_query_ref)
                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,
                    query_kpts,
                    ref_image,
                    ref_kpts,
                    matches[:50],
                    None,
                    flags=2,
                )

                match_file_name = f"match_image_{filename}.jpg"
                cv2.imwrite(os.path.join(demo_image_path , match_file_name), matched_image)

                # wrap_image = cv.warpPerspective(query_image_resize, M, (ref_image_resize.shape[1], ref_image_resize.shape[0]))
                wrap_image = cv.warpPerspective(query_image, M,
                                                (ref_image.shape[1], ref_image.shape[0]))
                # wrap_image = cv2.resize(wrap_image,(ref_image.shape[1], ref_image.shape[0]))
                wrap_filename = f"wrap_image_{filename}.jpg"
                cv2.imwrite(os.path.join(demo_image_path, wrap_filename), wrap_image)
                wraped_image_list.append((wrap_image,wrap_filename))

                result_image = cv2.subtract(ref_image, wrap_image)
                result_file_name = f"result_image_{filename}.jpg"
                cv2.imwrite(os.path.join(demo_image_path, result_file_name), result_image)
        except:
            pass

    return wraped_image_list

def read_params_from_yml(yml_file_path):
    with open(yml_file_path, 'r') as file:
        params = yaml.safe_load(file)
    return params



if __name__ == "__main__":
    yml_file_path = 'params.yml'
    params = read_params_from_yml(yml_file_path)

    registration_demo(params['image_dir'],params['demo_image_path'], params['json_ref_path'], params['ref_image_path'])