import hashlib
import json
import os.path
import math
from scipy.optimize import linear_sum_assignment
import numpy as np
from PIL import Image
import fastproot


# location to store array cache
CACHE_FILE_PATH = "./cache//point_array_cache.json"
LOADED_IMAGE_PATH = "./graphics/loaded.png"

CACHE_FILE_PATH = os.path.normpath(CACHE_FILE_PATH)
LOADED_IMAGE_PATH = os.path.normpath(LOADED_IMAGE_PATH)

class Point2D:
    x = 0
    y = 0
    color: tuple[int, int, int] = (0, 0, 0)

    def __init__(self, x, y, color: tuple[int, int, int] = (0, 0, 0)):
        self.x = x
        self.y = y
        self.color = color

    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
        new_color = tuple(int((1 - percentage) * self.color[i] + percentage * other.color[i]) for i in range(3))
        return Point2D(new_x, new_y, new_color)

    def __eq__(self, other):
        return (self.x, self.y) == (other.x, other.y)


def divide_points_into_groups(points):
    result = fastproot.divide_into_groups(points)
    res = [[Point2D(p[0], p[1], p[2]) for p in arr] for arr in result]
    return res

def _divide_points_into_groups(points):
    def flood_fill(point, group):
        if point not in group:
            group.append(point)

            neighbors = [(point.x + 1, point.y), (point.x - 1, point.y), (point.x, point.y + 1), (point.x, point.y - 1)]
            for neighbor_coords in neighbors:
                neighbor = next((p for p in points if p.x == neighbor_coords[0] and p.y == neighbor_coords[1]), None)
                if neighbor and neighbor not in group:
                    flood_fill(neighbor, group)

    groups = []
    remaining_points = [point for point in points if point.x < 64]  # Filter points with x < 64

    while remaining_points:
        group = []
        flood_fill(remaining_points[0], group)
        groups.append(group)

        # Remove points in the group from the remaining points
        remaining_points = [point for point in remaining_points if point not in group]

    return groups


def compute_bounding_box(points: list[Point2D]) -> tuple[float, float, float, float]:
    min_x = min(point.x for point in points)
    max_x = max(point.x for point in points)
    min_y = min(point.y for point in points)
    max_y = max(point.y for point in points)
    return min_x, max_x, min_y, max_y

def calculate_distance(points1: list[Point2D], points2: list[Point2D]) -> float:
    # Calculate the distance between the centers of two groups
    center1_x = (points1[0].x + points1[-1].x) / 2
    center1_y = (points1[0].y + points1[-1].y) / 2
    center2_x = (points2[0].x + points2[-1].x) / 2
    center2_y = (points2[0].y + points2[-1].y) / 2
    return ((center1_x - center2_x) ** 2 + (center1_y - center2_y) ** 2) ** 0.5

def pair_groups(set_a: list[list[Point2D]], set_b: list[list[Point2D]]) -> list[tuple[list[Point2D], list[Point2D]]]:
    pairs = []
    
    # Create dictionaries to store bounding boxes for each group
    bounding_boxes_a: dict[int, tuple[float, float, float, float]] = {}
    bounding_boxes_b: dict[int, tuple[float, float, float, float]] = {}

    # Calculate bounding boxes for all groups in both sets
    for i, group in enumerate(set_a + set_b):
        bounding_box = compute_bounding_box(group)
        if i < len(set_a):
            bounding_boxes_a[i] = bounding_box
        else:
            bounding_boxes_b[i - len(set_a)] = bounding_box

    # Check for overlaps and determine pairs
    for i, group_a in enumerate(set_a):
        overlap_detected = False
        for j, group_b in enumerate(set_b):
            bounding_box_a = bounding_boxes_a[i]
            bounding_box_b = bounding_boxes_b[j]
            if (
                bounding_box_a[0] <= bounding_box_b[1] and
                bounding_box_a[1] >= bounding_box_b[0] and
                bounding_box_a[2] <= bounding_box_b[3] and
                bounding_box_a[3] >= bounding_box_b[2]
            ):
                pairs.append((group_a, group_b))
                overlap_detected = True
                break
        
        if not overlap_detected:
            # Find the nearest neighbor in set B
            nearest_group = min(set_b, key=lambda group: calculate_distance(group_a, group))
            pairs.append((group_a, nearest_group))

    return pairs

def mirror_points(points: list[Point2D]) -> list[Point2D]:
    mirrored_points = []
    for point in points:
        mirrored_x = 128 - point.x  # Calculate the mirrored x-coordinate
        mirrored_point = Point2D(mirrored_x, point.y, point.color)
        mirrored_points.append(mirrored_point)
    return mirrored_points


def get_image_hash(image):
    image_hash = hashlib.sha1(image.tobytes()).hexdigest()
    return image_hash


def load_cached_point_arrays():
    cached_point_arrays = {}

    if os.path.isfile(CACHE_FILE_PATH):
        with open(CACHE_FILE_PATH, "r") as file:
            cached_point_arrays = json.load(file)

    return cached_point_arrays


def save_cached_point_arrays(cached_point_arrays):
    if not os.path.isfile(CACHE_FILE_PATH):
        open(CACHE_FILE_PATH, "x").close()

    with open(CACHE_FILE_PATH, "w") as file:
        json.dump(cached_point_arrays, file)


def generate_point_array_from_image(image):
    image.load()
    image = image.convert("RGB")  # Convert image to RGB color mode
    image_hash = get_image_hash(image)
    cached_point_arrays = load_cached_point_arrays()

    if image_hash in cached_point_arrays:
        print("Found existing point array matching png. Using existing array.")
        return [Point2D(point["x"], point["y"], tuple(point["color"])) for point in cached_point_arrays[image_hash]]

    print("No existing point array matching png found. Generating now.")
    width, height = image.size
    pixel_array = []

    image.save(LOADED_IMAGE_PATH)

    for y in range(height):
        for x in range(width):
            pixel = image.getpixel((x, y))
            if sum(pixel) > 15:  # any pixel which total color value is greater than 15.
                point = {"x": x, "y": y, "color": pixel}
                pixel_array.append(point)

    cached_point_arrays[image_hash] = pixel_array
    save_cached_point_arrays(cached_point_arrays)

    print("Point array generated and stored.")
    return [Point2D(point["x"], point["y"], tuple(point["color"])) for point in pixel_array]
    
    
def generate_image_from_point_array(points: list[Point2D], width: int, height: int) -> Image:
    # 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] = point.color
    
    return image


def interpolate_point_pairs(pairs: list[tuple[Point2D, Point2D]], percentage: float) -> list[Point2D]:
    interpolated_points:list[Point2D] = []
    for pair in pairs:
        point1, point2 = pair
        interpolated_point = point1.interpolate(point2, percentage)
        interpolated_points.append(interpolated_point)
    return interpolated_points


def pair_points(points1: list[Point2D], points2: list[Point2D]) -> list[tuple[Point2D, Point2D]]:
    # Update the size of the point arrays
    size1 = len(points1)
    size2 = len(points2)
    
    # 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
        
    row_ind, col_ind = fastproot.solve(points1, points2)

    # 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 _pair_points(points1: list[Point2D], points2: list[Point2D]) -> list[tuple[Point2D, Point2D]]:
    # Update the size of the point arrays
    size1 = len(points1)
    size2 = len(points2)
    
    # 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
    
    # 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