# See https://stackoverflow.com/a/39757388
from __future__ import annotations

from functools import lru_cache
from operator import attrgetter
from typing import TYPE_CHECKING, Iterable

from funcy import juxt, rpartial, compose, lflatten, first, second

from cv_analysis.utils import lift

if TYPE_CHECKING:
    from cv_analysis.utils.rectangle import Rectangle


def adjacent(alpha: Rectangle, beta: Rectangle, tolerance=7, strict=False):
    """Checks if the two rectangles are adjacent to each other.

    Args:
        alpha: The first rectangle.
        beta: The second rectangle.
        tolerance: The maximum distance between the two rectangles.
        strict: If True, the rectangles must be adjacent along one axis and contained within the other axis. Else, the
            rectangles must be adjacent along one axis and overlapping the other axis.
    Returns:
        True if the two rectangles are adjacent to each other, False otherwise.
    """
    select_strictness_variant = first if strict else second
    test_candidates = [
        # +---+
        # |   | +---+
        # | a | | b |
        # |   | +___+
        # +___+
        (right_left_aligned_and_vertically_contained, right_left_aligned_and_vertically_overlapping),
        #       +---+
        # +---+ |   |
        # | b | | a |
        # +___+ |   |
        #       +___+
        (left_right_aligned_and_vertically_contained, left_right_aligned_and_vertically_overlapping),
        # +-----------+
        # |     a     |
        # +___________+
        #    +-----+
        #    |  b  |
        #    +_____+
        (bottom_top_aligned_and_horizontally_contained, bottom_top_aligned_and_horizontally_overlapping),
        #    +-----+
        #    |  b  |
        #    +_____+
        # +-----------+
        # |     a     |
        # +___________+
        (top_bottom_aligned_and_horizontally_contained, top_bottom_aligned_and_horizontally_overlapping),
    ]

    tests = map(select_strictness_variant, test_candidates)
    return any(juxt(*tests)(alpha, beta, tolerance))


def right_left_aligned_and_vertically_overlapping(alpha: Rectangle, beta: Rectangle, tol):
    """Checks if the first rectangle is left of the other within a tolerance and also overlaps the other's y range."""
    return adjacent_along_one_axis_and_overlapping_along_perpendicular_axis(
        alpha.x2, beta.x1, beta.y1, beta.y2, alpha.y1, alpha.y2, tolerance=tol
    )


def left_right_aligned_and_vertically_overlapping(alpha: Rectangle, beta: Rectangle, tol):
    """Checks if the first rectangle is right of the other within a tolerance and also overlaps the other's y range."""
    return adjacent_along_one_axis_and_overlapping_along_perpendicular_axis(
        alpha.x1, beta.x2, beta.y1, beta.y2, alpha.y1, alpha.y2, tolerance=tol
    )


def bottom_top_aligned_and_horizontally_overlapping(alpha: Rectangle, beta: Rectangle, tol):
    """Checks if the first rectangle is above the other within a tolerance and also overlaps the other's x range."""
    return adjacent_along_one_axis_and_overlapping_along_perpendicular_axis(
        alpha.y2, beta.y1, beta.x1, beta.x2, alpha.x1, alpha.x2, tolerance=tol
    )


def top_bottom_aligned_and_horizontally_overlapping(alpha: Rectangle, beta: Rectangle, tol):
    """Checks if the first rectangle is below the other within a tolerance and also overlaps the other's x range."""
    return adjacent_along_one_axis_and_overlapping_along_perpendicular_axis(
        alpha.y1, beta.y2, beta.x1, beta.x2, alpha.x1, alpha.x2, tolerance=tol
    )


def right_left_aligned_and_vertically_contained(alpha: Rectangle, beta: Rectangle, tol):
    """Checks if the first rectangle is left of the other within a tolerance and also contains the other's y range."""
    return adjacent_along_one_axis_and_contained_within_perpendicular_axis(
        alpha.x2, beta.x1, beta.y1, beta.y2, alpha.y1, alpha.y2, tolerance=tol
    )


def left_right_aligned_and_vertically_contained(alpha: Rectangle, beta: Rectangle, tol):
    """Checks if the first rectangle is right of the other within a tolerance and also contains the other's y range."""
    return adjacent_along_one_axis_and_contained_within_perpendicular_axis(
        alpha.x1, beta.x2, beta.y1, beta.y2, alpha.y1, alpha.y2, tolerance=tol
    )


def bottom_top_aligned_and_horizontally_contained(alpha: Rectangle, beta: Rectangle, tol):
    """Checks if the first rectangle is above the other within a tolerance and also contains the other's x range."""
    return adjacent_along_one_axis_and_contained_within_perpendicular_axis(
        alpha.y2, beta.y1, beta.x1, beta.x2, alpha.x1, alpha.x2, tolerance=tol
    )


def top_bottom_aligned_and_horizontally_contained(alpha: Rectangle, beta: Rectangle, tol):
    """Checks if the first rectangle is below the other within a tolerance and also contains the other's x range."""
    return adjacent_along_one_axis_and_contained_within_perpendicular_axis(
        alpha.y1, beta.y2, beta.x1, beta.x2, alpha.x1, alpha.x2, tolerance=tol
    )


def adjacent_along_one_axis_and_overlapping_along_perpendicular_axis(
    axis_0_point_1,
    axis_1_point_2,
    axis_1_contained_point_1,
    axis_1_contained_point_2,
    axis_1_lower_bound,
    axis_1_upper_bound,
    tolerance,
):
    """Checks if two points are adjacent along one axis and two other points overlap a range along the perpendicular
    axis.
    """
    return adjacent_along_one_axis_and_overlapping_or_contained_along_perpendicular_axis(
        axis_0_point_1,
        axis_1_point_2,
        axis_1_contained_point_1,
        axis_1_contained_point_2,
        axis_1_lower_bound,
        axis_1_upper_bound,
        tolerance,
        mode="overlapping",
    )


def adjacent_along_one_axis_and_contained_within_perpendicular_axis(
    axis_0_point_1,
    axis_1_point_2,
    axis_1_contained_point_1,
    axis_1_contained_point_2,
    axis_1_lower_bound,
    axis_1_upper_bound,
    tolerance,
):
    """Checks if two points are adjacent along one axis and two other points overlap a range along the perpendicular
    axis.
    """
    return adjacent_along_one_axis_and_overlapping_or_contained_along_perpendicular_axis(
        axis_0_point_1,
        axis_1_point_2,
        axis_1_contained_point_1,
        axis_1_contained_point_2,
        axis_1_lower_bound,
        axis_1_upper_bound,
        tolerance,
        mode="contained",
    )


def adjacent_along_one_axis_and_overlapping_or_contained_along_perpendicular_axis(
    axis_0_point_1,
    axis_1_point_2,
    axis_1_contained_point_1,
    axis_1_contained_point_2,
    axis_1_lower_bound,
    axis_1_upper_bound,
    tolerance,
    mode,
):
    """Checks if two points are adjacent along one axis and two other points overlap a range along the perpendicular
    axis or are contained in that range, depending on the mode specified.
    """
    assert mode in ["overlapping", "contained"]
    quantifier = any if mode == "overlapping" else all
    return all(
        [
            abs(axis_0_point_1 - axis_1_point_2) <= tolerance,
            quantifier(
                [
                    axis_1_lower_bound <= p <= axis_1_upper_bound
                    for p in [axis_1_contained_point_1, axis_1_contained_point_2]
                ]
            ),
        ]
    )


def contains(alpha: Rectangle, beta: Rectangle, tol=3):
    """Checks if the first rectangle contains the second rectangle."""
    return (
        beta.x1 + tol >= alpha.x1
        and beta.y1 + tol >= alpha.y1
        and beta.x2 - tol <= alpha.x2
        and beta.y2 - tol <= alpha.y2
    )


def is_contained(rectangle: Rectangle, rectangles: Iterable[Rectangle]):
    """Checks if the rectangle is contained within any of the other rectangles."""
    other_rectangles = filter(lambda r: r != rectangle, rectangles)
    return any(map(rpartial(contains, rectangle), other_rectangles))


def iou(alpha: Rectangle, beta: Rectangle):
    """Calculates the intersection area over the union area of two rectangles."""
    return intersection(alpha, beta) / union(alpha, beta)


def area(rectangle: Rectangle):
    """Calculates the area of a rectangle."""
    return abs((rectangle.x2 - rectangle.x1) * (rectangle.y2 - rectangle.y1))


def union(alpha: Rectangle, beta: Rectangle):
    """Calculates the union area of two rectangles."""
    return area(alpha) + area(beta) - intersection(alpha, beta)


@lru_cache(maxsize=1000)
def intersection(alpha, beta):
    """Calculates the intersection of two rectangles."""
    return intersection_along_x_axis(alpha, beta) * intersection_along_y_axis(alpha, beta)


def intersection_along_x_axis(alpha, beta):
    """Calculates the intersection along the x-axis."""
    return intersection_along_axis(alpha, beta, "x")


def intersection_along_y_axis(alpha, beta):
    """Calculates the intersection along the y-axis."""
    return intersection_along_axis(alpha, beta, "y")


def intersection_along_axis(alpha, beta, axis):
    """Calculates the intersection along the given axis.

    Cases:
       a      b
    [-----] (---)  ==> [a1, b1, a2, b2] ==> max(0, (a2 - b1)) = 0
      b      a
    (---) [-----]  ==> [b1, a1, b2, a2] ==> max(0, (b2 - a1)) = 0
        a   b
    [--(----]----) ==> [a1, b1, a2, b2] ==> max(0, (a2 - b1)) = (a2 - b1)
        a b
    (-[---]----)   ==> [b1, a1, a2, b2] ==> max(0, (a2 - a1)) = (a2 - a1)
        b a
    [-(---)----]   ==> [a1, b1, b2, a2] ==> max(0, (b2 - b1)) = (b2 - b1)
        b    a
    (----[--)----] ==> [b1, a1, b2, a2] ==> max(0, (b2 - a1)) = (b2 - a1)
    """
    assert axis in ["x", "y"]

    def get_component_accessor(component):
        """Returns a function that accesses the given component of a rectangle."""
        return attrgetter(f"{axis}{component}")

    def make_access_components_and_sort_fn(component):
        """Returns a function that accesses and sorts the given component of multiple rectangles."""
        assert component in [1, 2]
        return compose(sorted, lift(get_component_accessor(component)))

    sort_first_components, sort_second_components = map(make_access_components_and_sort_fn, [1, 2])

    min_c1, max_c1, min_c2, max_c2 = lflatten(juxt(sort_first_components, sort_second_components)((alpha, beta)))
    intersection = max(0, min_c2 - max_c1)
    return intersection


def related(alpha: Rectangle, beta: Rectangle):
    return close(alpha, beta) or overlap(alpha, beta)


def close(alpha: Rectangle, beta: Rectangle, max_gap=14):
    # FIXME: Parameterize via factory
    return adjacent(alpha, beta, tolerance=max_gap, strict=True)


def overlap(alpha: Rectangle, beta: Rectangle):
    return intersection(alpha, beta) > 0