import os
from pathlib import Path
import cv2 as cv
import numpy as np

MIN_NUM_KEYPOINT_MATCHES = 50


def main():
    """
    Tworzy pętlę przechodzącą przez zawartość dwóch katalogów
    z odpowiadającymi sobie obrazami, rejestruje obrazy
    i wyświetla je naprzemiennie.
    """
    night1_files = sorted(os.listdir('night_1'))
    night2_files = sorted(os.listdir('night_2'))
    path1 = Path.cwd() / 'night_1'
    path2 = Path.cwd() / 'night_2'
    path3 = Path.cwd() / 'night_1_registered'

    for i, _ in enumerate(night1_files):
        img1 = cv.imread(str(path1 / night1_files[i]), cv.IMREAD_GRAYSCALE)
        img2 = cv.imread(str(path2 / night2_files[i]), cv.IMREAD_GRAYSCALE)

        print("Porównywanie {} z {}.\n"
              .format(night1_files[i], night2_files[i]))

        # Znajduje punkty kluczowe i najlepsze dopasowania między nimi.
        kp1, kp2, best_matches = find_best_matches(img1, img2)
        img_match = cv.drawMatches(img1, kp1, img2, kp2,
                                   best_matches, outImg=None)

        # Rysuje linię między dwoma obrazami.
        height, width = img1.shape
        cv.line(img_match, (width, 0), (width, height), (255, 255, 255), 1)
        QC_best_matches(img_match)  # Wykomentuj, aby pominąć ten krok.

        # Zarejestruj obraz po lewej stronie z użyciem punktów kluczowych.        
        img1_registered = register_image(img1, img2, kp1, kp2, best_matches)

        # Kontrola jakości rejestracji i zapisanie zarejestrowanego obrazu
        # (kroki opcjonalne):
        blink(img1, img1_registered, 'Sprawdzenie rejestracji', num_loops=5)
        out_filename = '{}_registered.png'.format(night1_files[i][:-4])
        # Spowoduje nadpisanie obrazu!
        cv.imwrite(str(path3 / out_filename), img1_registered)

        cv.destroyAllWindows()

        # Uruchamia komparator błyskowy.
        blink(img1_registered, img2, 'Komparator blyskowy', num_loops=15)


def find_best_matches(img1, img2):
    """
    Zwraca listę punktów kluczowych oraz listę najlepszych dopasowań
    między dwoma obrazami.
    """
    orb = cv.ORB_create(nfeatures=100)  # Tworzy obiekt ORB.

    # Znajduje punkty kluczowe i deskryptory z użyciem obiektu ORB.
    kp1, desc1 = orb.detectAndCompute(img1, mask=None)
    kp2, desc2 = orb.detectAndCompute(img2, mask=None)

    # Znajduje dopasowania punktów kluczowych z użyciem obiektu BMF.
    bf = cv.BFMatcher(cv.NORM_HAMMING, crossCheck=True)
    matches = bf.match(desc1, desc2)

    # Sortuje dopasowania w kolejności rosnącej
    # i zachowuje n najlepszych dopasowań.
    matches = sorted(matches, key=lambda x: x.distance)
    best_matches = matches[:MIN_NUM_KEYPOINT_MATCHES]

    return kp1, kp2, best_matches


def QC_best_matches(img_match):
    """Rysuje dopasowania punktów kluczowych połączone kolorowymi liniami."""
    cv.imshow('{} najlepiej dopasowanych punktow kluczowych'
              .format(MIN_NUM_KEYPOINT_MATCHES), img_match)
    cv.waitKey(2500)  # Okno będzie aktywne przez 2,5 s.


def register_image(img1, img2, kp1, kp2, best_matches):
    """Zwraca pierwszy obraz zarejestrowany do drugiego."""
    if len(best_matches) >= MIN_NUM_KEYPOINT_MATCHES:
        src_pts = np.zeros((len(best_matches), 2), dtype=np.float32)
        dst_pts = np.zeros((len(best_matches), 2), dtype=np.float32)
        for i, match in enumerate(best_matches):
            src_pts[i, :] = kp1[match.queryIdx].pt
            dst_pts[i, :] = kp2[match.trainIdx].pt
        h_array, mask = cv.findHomography(src_pts, dst_pts, cv.RANSAC)
        height, width = img2.shape  # Pozyskuje wymiary drugiego obrazu.
        img1_warped = cv.warpPerspective(img1, h_array, (width, height))

        return img1_warped

    else:
        print("OSTRZEŻENIE: Liczba dopasowań punktów kluczowych < {}\n"
              .format(MIN_NUM_KEYPOINT_MATCHES))
        return img1


def blink(image_1, image_2, window_name, num_loops):
    """
    Odwzorowuje działanie komparatora błyskowego
    wyświetlającego dwa obrazy.
    """
    for _ in range(num_loops):
        cv.imshow(window_name, image_1)
        cv.waitKey(330)
        cv.imshow(window_name, image_2)
        cv.waitKey(330)


if __name__ == '__main__':
    main()
