import numpy as np                             # numpy
import cv2 as cv                               # opencv 庫
import itertools as it                         # 迭代器
from multiprocessing.pool import ThreadPool    # 多進(jìn)程
from matplotlib import pyplot as plt           # 畫圖工具
import os


FLANN_INDEX_KDTREE = 0
FLANN_INDEX_LSH = 6

def affine_skew(tilt, phi, img, mask=None):
    h, w = img.shape[:2]
    if mask is None:
        mask = np.zeros((h, w), np.uint8)
        mask[:] = 255
    A = np.float32([[1, 0, 0], [0, 1, 0]])
    if phi != 0.0:
        phi = np.deg2rad(phi)
        s, c = np.sin(phi), np.cos(phi)
        A = np.float32([[c, -s], [s, c]])
        corners = [[0, 0], [w, 0], [w, h], [0, h]]
        tcorners = np.int32(np.dot(corners, A.T))
        x, y, w, h = cv.boundingRect(tcorners.reshape(1, -1, 2))
        A = np.hstack([A, [[-x], [-y]]])
        img = cv.warpAffine(img, A, (w, h), flags=cv.INTER_LINEAR, borderMode=cv.BORDER_REPLICATE)
    if tilt != 1.0:
        s = 0.8 * np.sqrt(tilt * tilt - 1)
        img = cv.GaussianBlur(img, (0, 0), sigmaX=s, sigmaY=0.01)
        img = cv.resize(img, (0, 0), fx=1.0 / tilt, fy=1.0, interpolation=cv.INTER_NEAREST)
        A[0] /= tilt
    if phi != 0.0 or tilt != 1.0:
        h, w = img.shape[:2]
        mask = cv.warpAffine(mask, A, (w, h), 
             flags=cv.INTER_NEAREST)
    ai = cv.invertAffineTransform(A)
    return img, mask, ai


def affine_detect(detector, img, mask=None, pool=None):
    params = [(1.0, 0.0)]
    for t in 2 ** (0.5 * np.arange(1, 6)):
        for phi in np.arange(0, 180, 72.0 / t):
            params.append((t, phi))

    def f(p):
        t, phi = p
        timg, tmask, Ai = affine_skew(t, phi, img)
        keypoints, descrs = detector.detectAndCompute(timg, tmask)
        for kp in keypoints:
            x, y = kp.pt
            kp.pt = tuple(np.dot(Ai, (x, y, 1)))
        if descrs is None:
            descrs = []
        return keypoints, descrs

    keypoints, descrs = [], []
    if pool is None:
        ires = it.imap(f, params)
    else:
        ires = pool.imap(f, params)



    for i, (k, d) in enumerate(ires):
        # print('affine sampling: %d / %d\r' % (i + 1, len(params)), end='')
        keypoints.extend(k)
        descrs.extend(d)

    return keypoints, np.array(descrs)


#初始化算法
def init_feature(name):
    chunks = name.split('-')
    if chunks[0] == 'sift':
        detector = cv.xfeatures2d.SIFT_create()
        norm = cv.NORM_L2
    elif chunks[0] == 'surf':
        detector = cv.xfeatures2d.SURF_create(800)
        norm = cv.NORM_L2
    elif chunks[0] == 'orb':
        detector = cv.ORB_create(800)
        norm = cv.NORM_HAMMING
    elif chunks[0] == 'akaze':
        detector = cv.AKAZE_create()
        norm = cv.NORM_HAMMING
    elif chunks[0] == 'brisk':
        detector = cv.BRISK_create()
        norm = cv.NORM_HAMMING
    else:
        return None, None
    if 'flann' in chunks:
        if norm == cv.NORM_L2:
            flann_params = dict(algorithm=FLANN_INDEX_KDTREE, trees=5)
        else:
            flann_params = dict(algorithm=FLANN_INDEX_LSH,
                                table_number=6,  # 12
                                key_size=12,  # 20
                                multi_probe_level=1)  # 2
   

def main(img1, img2):
    feature_name = 'orb-flann'
    detector, matcher = init_feature(feature_name)

    pool = ThreadPool(processes=cv.getNumberOfCPUs())

    kp1, desc1 = affine_detect(detector, img1, pool=pool)
    kp2, desc2 = affine_detect(detector, img2, pool=pool)

    matches = matcher.knnMatch(desc1, desc2, k=2)

    # lowe‘s過濾
    good = []
    for m, n in matches:
        if m.distance < 0.7 * n.distance:
            good.append(m)
    ratio1 = len(good) / len(matches)
    print('匹配率1', ratio1)
    # ransac 隨機抽樣一直過濾
    min_match_count = 10  # ransac閾值，期望匹配的好點最低數(shù)目 (根據(jù)項目判斷)
    inLiner = 0  # 符合ransac算法變換矩陣的好點

    if len(good) > min_match_count:

        src_pts = np.float32([kp1[m.queryIdx].pt for m in good]).reshape(-1, 1, 2)
        dst_pts = np.float32([kp2[m.trainIdx].pt for m in good]).reshape(-1, 1, 2)

        # M: 3x3 變換矩陣.
        M, mask = cv.findHomography(src_pts, dst_pts, cv.RANSAC, 10.0)
        matchesMask = mask.ravel().tolist()
         for i in range(len(matchesMask)):
            if matchesMask[i] == 1:
                inLiner += 1

        ratio2 = inLiner / len(matchesMask)
        # return ratio1, ratio2, len(matches), len(good), inLiner
        print('匹配率2',ratio2)
        draw_params = dict(matchColor=(0, 255, 0),
                           singlePointColor=(0, 0, 255),
                           matchesMask=matchesMask,
                           flags=0)

        img3 = cv.drawMatches(img1, kp1, img2, kp2, good, None, **draw_params)
        plt.imshow(img3), plt.show()
    else:
        # 沒有足夠多的好點可以認(rèn)為直接不匹配
        print('Done')
    # return ratio1, 0, len(matches), len(good), inLiner


if __name__ == '__main__':
     img1 = cv.imread('original/12.jpg', 0)
     img2 = cv.imread('target/12.jpg', 0) 
     main(img2, img1)