import math
from PIL import Image
import numpy as np
from scipy.optimize import linear_sum_assignment

    
class Point2D:
    x = 0
    y = 0
    
    def __init__(self, x, y):
        self.x = x
        self.y = y
        
    def round(self):
        self.x = round(self.x)
        self.y = round(self.y)
        return self
    
    def distance(self, other):
        dx = self.x - other.x
        dy = self.y - other.y
        return math.sqrt(dx**2 + dy**2)
    
    def interpolate(self, other, percentage):
        new_x = self.x + (other.x - self.x) * percentage
        new_y = self.y + (other.y - self.y) * percentage
        return Point2D(new_x, new_y)
    
    def __eq__(self, other):
        return (self.x, self.y) == (other.x, other.y)


def generate_point_array_from_image(image):
    image = image.convert("RGB")  # Convert image to RGB color mode

    width, height = image.size
    point_array = []

    # Iterate over the pixels and generate Point2D instances
    for y in range(height):
        for x in range(width):
            pixel = image.getpixel((x, y))
            if pixel != (0, 0, 0):  # Assuming white pixels
                point = Point2D(x, y)
                point_array.append(point)

    return point_array


def generate_image_from_point_array(points, width, height):
    # Create a new blank image
    image = Image.new("RGB", (width, height), "black")
    
    # Set the pixels corresponding to the points as white
    pixels = image.load()
    for point in points:
        point = point.round()
        x = point.x
        y = point.y
        pixels[x, y] = (255, 255, 255)
    
    return image


def pair_points(points1, points2):
    # Determine the size of the point arrays
    size1 = len(points1)
    size2 = len(points2)
    
    # Create a cost matrix based on the distances between points
    cost_matrix = np.zeros((size1, size2))
    for i in range(size1):
        for j in range(size2):
            cost_matrix[i, j] = points1[i].distance(points2[j])
    
    # Duplicate points in the smaller array to match the size of the larger array
    if size1 > size2:
        num_duplicates = size1 - size2
        duplicated_points = np.random.choice(points2, size=num_duplicates).tolist()
        points2 += duplicated_points
    elif size2 > size1:
        num_duplicates = size2 - size1
        duplicated_points = np.random.choice(points1, size=num_duplicates).tolist()
        points1 += duplicated_points
    
    # Update the size of the point arrays
    size1 = len(points1)
    size2 = len(points2)
    
    # Create a new cost matrix with the updated sizes
    cost_matrix = np.zeros((size1, size2))
    for i in range(size1):
        for j in range(size2):
            cost_matrix[i, j] = points1[i].distance(points2[j])
    
    # Solve the assignment problem using the Hungarian algorithm
    row_ind, col_ind = linear_sum_assignment(cost_matrix)
    
    # Create pairs of points based on the optimal assignment
    pairs = []
    for i, j in zip(row_ind, col_ind):
        pairs.append((points1[i], points2[j]))
    
    return pairs


def interpolate_point_pairs(pairs, percentage):
    interpolated_points = []
    for pair in pairs:
        point1, point2 = pair
        interpolated_point = point1.interpolate(point2, percentage)
        interpolated_points.append(interpolated_point)
    return interpolated_points


Image1 = Image.open("CiscoTheProot/faces/prootface3.png")
Image2 = Image.open("CiscoTheProot/faces/prootface4.png")

pixelArray1 = generate_point_array_from_image(Image1)
pixelArray2 = generate_point_array_from_image(Image2)

pairs = pair_points(pixelArray1, pixelArray2)


generate_image_from_point_array(interpolate_point_pairs(pairs, 0), 128, 32).show()
generate_image_from_point_array(interpolate_point_pairs(pairs, .25), 128, 32).show()
generate_image_from_point_array(interpolate_point_pairs(pairs, .5), 128, 32).show()
generate_image_from_point_array(interpolate_point_pairs(pairs, .75), 128, 32).show()
generate_image_from_point_array(interpolate_point_pairs(pairs, 1), 128, 32).show()

print(pairs)