import numpy as np
import numpy.typing as npt
from matplotlib import pyplot as plt
import cv2
import uz_framework.image as uz_image
import uz_framework.text as uz_text
import os

##############################################
# EXCERCISE 1: Exercise 1: Image derivatives #
##############################################

def ex1():
	one_a()
	one_b()

def one_a() -> None:
	img = uz_image.imread_gray("data/graf/graf_a.jpg", uz_image.ImageType.float64)
	sigmas = [3, 6, 9, 12]
	# Plot the points
	fig, axs = plt.subplots(2, len(sigmas))
	fig.suptitle("Hessian corner detection")

	for i, sigma in enumerate(sigmas):
		determinant, hessian_points  = uz_image.hessian_points(img, sigma, 0.004)
		# Plot determinant
		axs[0, i].imshow(determinant)
		axs[0, i].set_title(f"Sigma: {sigma}")
		# Plot grayscale image
		axs[1, i].imshow(img, cmap="gray")
		# Plot scatter hessian points (x, y)
		axs[1, i].scatter(hessian_points[:, 1], hessian_points[:, 0], s=20, c="r", marker="x")

	plt.show()

def one_b() -> None:
	img = uz_image.imread_gray("data/graf/graf_a.jpg", uz_image.ImageType.float64)
	sigmas = [3, 6, 9]
	# Plot the points
	fig, axs = plt.subplots(2, len(sigmas))
	fig.suptitle("Harris corner detection")

	for i, sigma in enumerate(sigmas):
		determinant, harris_points  = uz_image.harris_detector(img, sigma, treshold=1e-6)
		# Plot determinant
		axs[0, i].imshow(determinant)
		axs[0, i].set_title(f"Sigma: {sigma}")
		# Plot grayscale image
		axs[1, i].imshow(img, cmap="gray")
		# Plot scatter hessian points
		axs[1, i].scatter(harris_points[:, 1], harris_points[:, 0], s=20, c="r", marker="x")

	plt.show()

def ex2():
	two_a()

def two_a() -> None:
	"""
	Hello
	"""
	graph_a_small = uz_image.imread_gray("data/graf/graf_a_small.jpg", uz_image.ImageType.float64)
	graph_b_small = uz_image.imread_gray("data/graf/graf_b_small.jpg", uz_image.ImageType.float64)

	# Get the keypoints
	_, graph_a_keypoints = uz_image.harris_detector(graph_a_small, 3, treshold=1e-6)
	_, graph_b_keypoints = uz_image.harris_detector(graph_b_small, 3, treshold=1e-6)

	# Get the descriptors
	graph_a_descriptors = uz_image.simple_descriptors(graph_a_small, graph_a_keypoints[:,0], graph_a_keypoints[:,1])
	graph_b_descriptors = uz_image.simple_descriptors(graph_b_small, graph_b_keypoints[:,0], graph_b_keypoints[:,1])

	# Find the correspondences
	matches_a = uz_image.find_correspondences(graph_a_descriptors, graph_b_descriptors)
	matches_b = uz_image.find_correspondences(graph_b_descriptors, graph_a_descriptors)

	matches_a_coordinates = []
	matche_b_coordinates = []
	for i, match in enumerate(matches_a):
		if i % 2 == 0: # plot every second one
			if np.flip(match) in matches_b: # Check if the match is reciprocal
				print(match)
				print(np.flip(match))
				print(matches_b)
				print(np.argwhere(matches_b == np.flip(match)))
				matches_a_coordinates.append(np.flip(graph_a_keypoints[match[0]]))
				matche_b_coordinates.append(np.flip(graph_b_keypoints[match[1]]))
			else:
				print("Not reciprocal")

	# Plot the matches
	uz_image.display_matches(graph_a_small, matches_a_coordinates, graph_b_small, matche_b_coordinates)

def two_b() -> None:
	"""
	jjjjj
	"""



	
# ######## #
# SOLUTION #
# ######## #

def main():
	#ex1() 
	ex2()

if __name__ == '__main__':
	main()