from __future__ import annotations

from collections.abc import Callable
from copy import deepcopy
from math import floor, inf, isnan, nan, sqrt
from typing import Generic

from ..array import argpass
from ..types import Accessor, Point2D, T
from .add import add
from .quad import Quad


def default_x(d: Point2D) -> float:
    return d[0]


def default_y(d: Point2D) -> float:
    return d[1]


class Quadtree(Generic[T]):
    """
    A quadtree recursively partitions two-dimensional space into squares,
    dividing each square into four equally-sized squares. Each distinct point
    exists in a unique leaf node; coincident points are represented by a linked
    list. Quadtrees can accelerate various spatial operations, such as the
    `Barnes–Hut approximation
    <https://en.wikipedia.org/wiki/Barnes%E2%80%93Hut_simulation>`_ for
    computing many-body forces, collision detection, and searching for nearby
    points.

    Parameters
    ----------
    x : Accessor[T, float]
        X accessor
    y : Accessor[T, float]
        Y accessor
    x0 : float
        Inclusive lower bound :math:`x_0`
    y0 : float
        Inclusive lower bound :math:`y_0`
    x1 : float
        Inclusive upper bound :math:`x_1`
    y1 : float
        Inclusive upper bound :math:`y_1`
    """

    def __init__(
        self,
        x: Accessor[T, float],
        y: Accessor[T, float],
        x0: float,
        y0: float,
        x1: float,
        y1: float,
    ):
        self._x = argpass(x)
        self._y = argpass(y)
        self._x0 = x0
        self._y0 = y0
        self._x1 = x1
        self._y1 = y1
        self._root = None

    def x(self, x: Accessor[T, float]) -> Quadtree:
        """
        Sets the x-coordinate accessor and returns the quadtree.
        The x accessor is used to derive the x coordinate of data when adding
        to and removing from the tree. It is also used when finding to
        re-access the coordinates of data previously added to the tree;
        therefore, the x and y accessors must be consistent, returning the same
        value given the same input.

        Parameters
        ----------
        x : Accessor[T, float]
            X accessor

        Returns
        -------
        Quadtree
            Itself
        """
        self._x = argpass(x)
        return self

    def y(self, y: Accessor[T, float]) -> Quadtree:
        """
        Sets the y-coordinate accessor and returns the quadtree.
        The y accessor is used to derive the y coordinate of data when adding
        to and removing from the tree. It is also used when finding to
        re-access the coordinates of data previously added to the tree;
        therefore, the x and y accessors must be consistent, returning the same
        value given the same input.

        Parameters
        ----------
        y : Accessor[T, float]
            Y accessor

        Returns
        -------
        Quadtree
            Itself
        """
        self._y = argpass(y)
        return self

    def add(self, d: T) -> Quadtree:
        """
        Adds the specified datum to the quadtree, deriving its coordinates
        :math:`(x, y)` using the current x and y accessors, and returns the
        quadtree.

        If the new point is outside the current extent of the quadtree, the
        quadtree is automatically expanded to cover the new point.

        Parameters
        ----------
        d : T
            Datum

        Returns
        -------
        Quadtree
            Itself
        """
        x = self._x(d)
        y = self._y(d)
        return add(self.cover(x, y), x, y, d)

    def add_all(self, data: list[T]) -> Quadtree:
        """
        Adds the specified iterable of data to the quadtree, deriving each
        element’s coordinates :math:`(x, y)` using the current x and y
        accessors, and return this quadtree.


        Parameters
        ----------
        data : list[T]
            List of data values

        Returns
        -------
        Quadtree
            Itself
        """
        data = list(data)
        n = len(data)
        xz = [None] * n
        yz = [None] * n
        x0 = inf
        y0 = x0
        x1 = -x0
        y1 = x1

        for i in range(n):
            d = data[i]
            x = self._x(d)
            y = self._y(d)
            if isnan(x) or isnan(y):
                continue
            xz[i] = x
            yz[i] = y
            if x < x0:
                x0 = x
            if x > x1:
                x1 = x
            if y < y0:
                y0 = y
            if y > y1:
                y1 = y

        if x0 > x1 or y0 > y1:
            return self

        self.cover(x0, y0).cover(x1, y1)

        for i in range(n):
            add(self, xz[i], yz[i], data[i])

        return self

    def cover(self, x: float, y: float) -> Quadtree:
        """
        Expands the quadtree to cover the specified point :math:`(x, y)`, and
        returns the quadtree.

        Parameters
        ----------
        x : float
            X value
        y : float
            Y value

        Returns
        -------
        Quadtree
            Itself
        """
        if isnan(x) or isnan(y):
            return self

        x0 = self._x0
        y0 = self._y0
        x1 = self._x1
        y1 = self._y1

        if isnan(x0):
            x0 = floor(x)
            y0 = floor(y)
            x1 = x0 + 1
            y1 = y0 + 1
        else:
            z = (x1 - x0) or 1
            node = self._root

            while x0 > x or x >= x1 or y0 > y or y >= y1:
                i = (y < y0) << 1 | (x < x0)
                parent = [None] * 4
                parent[i] = node
                node = parent
                z *= 2
                match i:
                    case 0:
                        x1 = x0 + z
                        y1 = y0 + z
                    case 1:
                        x0 = x1 - z
                        y1 = y0 + z
                    case 2:
                        x1 = x0 + z
                        y0 = y1 - z
                    case 3:
                        x0 = x1 - z
                        y0 = y1 - z

            if self._root and isinstance(self._root, list):
                self._root = node

        self._x0 = x0
        self._y0 = y0
        self._x1 = x1
        self._y1 = y1
        return self

    def set_extent(
        self, extent: list[tuple[float, float], tuple[float, float]]
    ) -> Quadtree:
        """
        Expands the quadtree to cover the specified points :code:`[[x0, y0],
        [x1, y1]]` and returns the quadtree.

        Parameters
        ----------
        extent : list[tuple[float, float], tuple[float, float]]
            Extent values

        Returns
        -------
        Quadtree
            Itself
        """
        return self.cover(extent[0][0], extent[0][1]).cover(extent[1][0], extent[1][1])

    def data(self) -> list[T]:
        """
        Returns an array of all data in the quadtree.

        Returns
        -------
        list[T]
            List of all data
        """
        data = []

        def visit(node):
            if isinstance(node, list):
                return
            while True:
                data.append(node["data"])
                node = node.get("next")
                if node is None:
                    break

        self.visit(visit)
        return data

    def find(self, x: float, y: float, radius: float | None = None) -> Quadtree:
        """
        Returns the datum closest to the position :math:`(x, y)` with the given
        search radius. If radius is not specified, it defaults to infinity.

        Parameters
        ----------
        x : float
            x value
        y : float
            y value
        radius : float | None
            radius value

        Returns
        -------
        Quadtree
            Itself
        """
        x0 = self._x0
        y0 = self._y0
        x3 = self._x1
        y3 = self._y1
        quads = []
        data = None
        node = self._root

        if node:
            quads.append(Quad(node, x0, y0, x3, y3))
        if radius is None:
            radius = inf
        else:
            x0 = x - radius
            y0 = y - radius
            x3 = x + radius
            y3 = y + radius
            radius *= radius

        while len(quads):
            q = quads.pop()
            node = q.node
            x1 = q.x0
            y1 = q.y0
            x2 = q.x1
            y2 = q.y1
            if not node or x1 > x3 or y1 > y3 or x2 < x0 or y2 < y0:
                continue

            if isinstance(node, list):
                xm = (x1 + x2) * 0.5
                ym = (y1 + y2) * 0.5

                quads.extend(
                    (
                        Quad(node[3], xm, ym, x2, y2),
                        Quad(node[2], x1, ym, xm, y2),
                        Quad(node[1], xm, y1, x2, ym),
                        Quad(node[0], x1, y1, xm, ym),
                    )
                )

                i = (y >= ym) << 1 | (x >= xm)
                if i:
                    q = quads[-1]
                    quads[-1] = quads[-1 - i]
                    quads[-1 - i] = q
            else:
                dx = x - self._x(node["data"])
                dy = y - self._y(node["data"])
                d2 = dx * dx + dy * dy
                if d2 < radius:
                    radius = d2
                    d = sqrt(d2)
                    x0 = x - d
                    y0 = y - d
                    x3 = x + d
                    y3 = y + d
                    data = node["data"]

        return data

    def remove(self, d: T) -> Quadtree:
        """
        Removes the specified datum from the quadtree, deriving its coordinates
        :math:`(x,y)` using the current x and y accessors, and returns the quadtree.

        Parameters
        ----------
        d : T
            Datum

        Returns
        -------
        Quadtree
            Itself
        """
        x = self._x(d)
        y = self._y(d)
        if isnan(x) or isnan(y):
            return self

        parent = None
        previous = None
        retainer = None
        node = self._root
        x0 = self._x0
        y0 = self._y0
        x1 = self._x1
        y1 = self._y1

        if node is None:
            return self

        if isinstance(node, list):
            while True:
                xm = (x0 + x1) * 0.5
                right = x >= xm
                if right:
                    x0 = xm
                else:
                    x1 = xm

                ym = (y0 + y1) * 0.5
                bottom = y >= ym
                if bottom:
                    y0 = ym
                else:
                    y1 = ym

                parent = node
                i = bottom << 1 | right
                node = node[i]
                if node is None:
                    return self
                if isinstance(node, dict):
                    break
                if parent[(i + 1) & 3] or parent[(i + 2) & 3] or parent[(i + 3) & 3]:
                    retainer = parent
                    j = i

        while node["data"] is not d:
            previous = node
            node = node.get("next")
            if node is None:
                return self
        next = node.pop("next", None)

        if previous:
            if next:
                previous["next"] = next
            else:
                previous.pop("next")
            return self

        if parent is None:
            self._root = next
            return self

        parent[i] = next

        node = parent[0] or parent[1] or parent[2] or parent[3]
        if node == (parent[3] or parent[2] or parent[1] or parent[0]) and isinstance(
            node, dict
        ):
            if retainer:
                retainer[j] = node
            else:
                self._root = node

        return self

    def remove_all(self, data: list[T]) -> Quadtree:
        """
        Removes the specified data from the quadtree, deriving their
        coordinates :math:`(x,y)` using the current x and y accessors, and
        returns the quadtree.

        Parameters
        ----------
        data : list[T]
            List of data

        Returns
        -------
        Quadtree
            Itself
        """
        for d in data:
            self.remove(d)
        return self

    def visit(
        self, callback: Callable[[list | dict | None, float, float, float, float], bool]
    ) -> Quadtree:
        """
        Visits each node in the quadtree in pre-order traversal, invoking the
        specified callback with arguments :code:`node`, :code:`x0`, :code:`y0`,
        :code:`x1`, :code:`y1` for each node, where node is the node being
        visited, :math:`(x_0, y_0)` are the lower bounds of the node, and
        :math:`(x_1, y_1)` are the upper bounds, and returns the quadtree.
        (Assuming that positive x is right and positive y is down, as is
        typically the case in Canvas and SVG, :math:`(x_0, y_0)` is the
        top-left corner and :math:`(x_1, y_1)` is the lower-right corner;
        however, the coordinate system is arbitrary, so more formally
        :math:`x_0 \\le x_1` and :math:`y0 \\le y1`.)

        If the callback returns true for a given node, then the children of
        that node are not visited; otherwise, all child nodes are visited. This
        can be used to quickly visit only parts of the tree, for example when
        using the `Barnes–Hut approximation
        <https://en.wikipedia.org/wiki/Barnes%E2%80%93Hut_simulation>`_. Note,
        however, that child quadrants are always visited in sibling order:
        top-left, top-right, bottom-left, bottom-right. In cases such as
        search, visiting siblings in a specific order may be faster.

        Parameters
        ----------
        callback : Callable[[list | dict | None, float, float, float, float], bool]
            Callback function

        Returns
        -------
        Quadtree
            Itself

        Notes
        -----
        The number of arguments for the :code:`callback` function is
        automatically determined. Therefore, you can pass 0 to 5 arguments to
        the :code:`callback` function. For example, `def callback(node):` is a
        valid function.
        """
        callback = argpass(callback)
        quads = []
        node = self._root
        if node:
            quads.append(Quad(node, self._x0, self._y0, self._x1, self._y1))
        while quads:
            q = quads.pop()
            node = q.node
            x0 = q.x0
            y0 = q.y0
            x1 = q.x1
            y1 = q.y1
            if not callback(node, x0, y0, x1, y1) and isinstance(node, list):
                xm = (x0 + x1) * 0.5
                ym = (y0 + y1) * 0.5
                child = node[3]
                if child:
                    quads.append(Quad(child, xm, ym, x1, y1))
                child = node[2]
                if child:
                    quads.append(Quad(child, x0, ym, xm, y1))
                child = node[1]
                if child:
                    quads.append(Quad(child, xm, y0, x1, ym))
                child = node[0]
                if child:
                    quads.append(Quad(child, x0, y0, xm, ym))
        return self

    def visit_after(
        self, callback: Callable[[dict, float, float, float, float], bool]
    ) -> Quadtree:
        """
        Visits each node in the quadtree in post-order traversal, invoking the
        specified callback with arguments :code:`node`, :code:`x0`, :code:`y0`,
        :code:`x1`, :code:`y1` for each node, where node is the node being
        visited, :code:`(x_0, y_0)` are the lower bounds of the node, and
        :code:`(x_1, y_1)` are the upper bounds, and returns the quadtree.
        (Assuming that positive :code:`x` is right and positive :code:`y` is
        down, as is typically the case in Canvas and SVG, :code:`(x_0, y_0)` is
        the top-left corner and :code:`(x_1, y_1)` is the lower-right corner;
        however, the coordinate system is arbitrary, so more formally
        :math:`x_0 \\lt x_1` and :math:`y0 \\lt y1`.). Returns root.

        Parameters
        ----------
        callback : Callable[[dict, float, float, float, float], bool]
            Callback function

        Returns
        -------
        Quadtree
            Itself

        Notes
        -----
        The number of arguments for the :code:`callback` function is
        automatically determined. Therefore, you can pass 0 to 5 arguments to
        the :code:`callback` function. For example, `def callback(node):` is a
        valid function.
        """
        callback = argpass(callback)
        quads = []
        next = []
        if self._root:
            quads.append(Quad(self._root, self._x0, self._y0, self._x1, self._y1))
        while quads:
            q = quads.pop()
            node = q.node
            if isinstance(node, list):
                x0 = q.x0
                y0 = q.y0
                x1 = q.x1
                y1 = q.y1
                xm = (x0 + x1) * 0.5
                ym = (y0 + y1) * 0.5
                child = node[0]
                if child:
                    quads.append(Quad(child, x0, y0, xm, ym))
                child = node[1]
                if child:
                    quads.append(Quad(child, xm, y0, x1, ym))
                child = node[2]
                if child:
                    quads.append(Quad(child, x0, ym, xm, y1))
                child = node[3]
                if child:
                    quads.append(Quad(child, xm, ym, x1, y1))
            next.append(q)
        while len(next) != 0:
            q = next.pop()
            callback(q.node, q.x0, q.y0, q.x1, q.y1)
        return self

    def copy(self) -> Quadtree:
        """
        Returns a deep copy of itself.

        Returns
        -------
        Quadtree
            Deep copy of itself
        """
        return deepcopy(self)

    def size(self) -> float:
        """
        Returns the total number of data in the quadtree.

        Returns
        -------
        float
            Total number of data in sthe quadtree.
        """
        size = [0]

        def visit(node: list | dict | None):
            if isinstance(node, list):
                return
            while True:
                size[0] += 1
                node = node.get("next")
                if node is None:
                    break

        self.visit(visit)
        return size[0]

    def get_extent(self) -> list[tuple[float, float], tuple[float, float]]:
        if isnan(self._x0):
            return None
        return [[self._x0, self._y0], [self._x1, self._y1]]

    def get_root(self) -> list | dict | None:
        return self._root

    def get_x(self) -> Accessor[T, float]:
        return self._x

    def get_y(self) -> Accessor[T, float]:
        return self._y


def quadtree(
    data: list[T] | None = None,
    x: Accessor[T, float] = None,
    y: Accessor[T, float] = None,
) -> Quadtree[T]:
    """
    Creates a new, empty quadtree with an empty extent and the default x and y
    accessors. If data is specified, adds the specified iterable of data to the
    quadtree.

    Parameters
    ----------
    data : list[T] | None
        List of data values
    x : Accessor[T, float]
        X accessor
    y : Accessor[T, float]
        Y accessor

    Returns
    -------
    Quadtree[T]
        Quadtree object
    """
    tree = Quadtree(
        default_x if x is None else x, default_y if y is None else y, nan, nan, nan, nan
    )
    return tree if data is None else tree.add_all(data)
