from math import sqrt, floor
from rectangles import Rect

def get_angle(rect):
    """Get angle of rotation in [0, 90].

    Args:
        rect (Rect): Named tuple Rect with rectangle parameters.
            Rect = namedtuple('Rect', ['x', 'y', 'width', 'height', 'angle'])

    Returns:
        float: Angle of rotation in [0, 90].
    """
    return rect.angle-floor(rect.angle/90.0)*90


def set_angle(rect, angle):
    """Change angle of rotation of rectangle.

    Args:
        rect (Rect): Named tuple Rect with rectangle parameters.
            Rect = namedtuple('Rect', ['x', 'y', 'width', 'height', 'angle'])
        angle (float): New angle of rotation.

    Return:
        (list): List of rectangle parameters.
    """
    return Rect(rect.x, rect.y, rect.width, rect.height, angle)


def distance(rect1, rect2):
    """Calculate distance between centers of rectangles.

    Args:
        rect1 (Rect): Named tuple Rect with rectangle parameters.
            Rect = namedtuple('Rect', ['x', 'y', 'width', 'height', 'angle'])
        rect2 (Rect): Named tuple Rect with rectangle parameters.
            Rect = namedtuple('Rect', ['x', 'y', 'width', 'height', 'angle'])

    Returns:
        float: Euclidean distance between centers of rectangles
    """
    return sqrt((rect1.x-rect2.x)**2+(rect1.y-rect2.y)**2)


def kneighbors(rects, n_neighbors, angle_delta, dist_delta):
    """Find k nearest rectangles for each rectangle.

    Args:
        rects (list): List of named tuples Rect with rectangles parameters.
            Rect = namedtuple('Rect', ['x', 'y', 'width', 'height', 'angle'])
        n_neighbors (int): Maximum number of neighbors.
        angle_delta (float): Maximum difference between angles of rotation of neighbors.
        dist_delta (float): Maximum distance between centers of neighbors.

    Returns:
        list: List of indices of neighbors and distances to them
            in the following format: [[[idx_1, dist_1], ..., [idx_k, dist_k]]]
    """
    neighbors = []
    for i in range(len(rects)):
        rect_neighbors = []
        for j in range(len(rects)):
            if i == j:
                continue
            angle_diff = abs(get_angle(rects[i])-get_angle(rects[j]))
            if angle_diff > angle_delta and 90-angle_diff > angle_delta:
                continue

            dist = distance(rects[i], rects[j])
            if dist > dist_delta:
                continue
            if len(rect_neighbors) < n_neighbors:
                rect_neighbors.append((j, dist))
            else:
                for k in range(n_neighbors):
                    if dist < rect_neighbors[k][1]:
                        rect_neighbors[k] = (j, dist)
                        break
        neighbors.append(rect_neighbors)
    return neighbors


def average_angles(rect, neighbors, angle_delta):
    """Average angles of rectangle and its neighbors.

    Args:
        rect (Rect): Named tuple Rect with rectangle parameters.
            Rect = namedtuple('Rect', ['x', 'y', 'width', 'height', 'angle'])
        neighbors (list): List of named tuples Rect with rectangles parameters.
        angle_delta (float): Maximum difference between angles of rotation of neighbors.

    Returns:
        float: average angle
    """
    neighbors_angles = [get_angle(neighbor) for neighbor in neighbors]
    for i in range(len(neighbors_angles)):
        if rect.angle > neighbors_angles[i]:
            while(rect.angle - neighbors_angles[i] > angle_delta):
                neighbors_angles[i] += 90
        else:
            while(neighbors_angles[i] - rect.angle > angle_delta):
                neighbors_angles[i] -= 90
    avg_angle = (rect.angle+sum(neighbors_angles))/(len(neighbors_angles)+1)
    return avg_angle


def align_angles(rects, n_neighbors, angle_delta, dist_delta, iterations):
    """Align angles between rectangles.
        Find k nearest rectangles for each rectangle and average their angles

    Args:
        rects (list): List of named tuples Rect with rectangles parameters.
            Rect = namedtuple('Rect', ['x', 'y', 'width', 'height', 'angle'])
        n_neighbors (int): Maximum number of neighbors.
        angle_delta (float): Maximum difference between angles of rotation of neighbors.
        dist_delta (float): Maximum distance between centers of neighbors.
        iterations (int): Number of iterations.

    Returns:
        list: List of named tuples Rect with aligned rectangles parameters.
            Rect = namedtuple('Rect', ['x', 'y', 'width', 'height', 'angle'])
    """
    processed_rects = rects[:]
    for iter_num in range(iterations):
        neighbors = kneighbors(processed_rects, n_neighbors, angle_delta, dist_delta)
        iter_rects = []
        for i in range(len(processed_rects)):
            neighbors_idx = [item[0] for item in neighbors[i]]
            rect_neighbors = [processed_rects[j] for j in neighbors_idx]
            new_angle = average_angles(processed_rects[i], rect_neighbors, angle_delta)
            iter_rects.append(set_angle(processed_rects[i], new_angle))
        processed_rects = iter_rects
    return processed_rects
