import cython
cimport cython
from geometry import MultipointBase

cdef int MAXCHILDREN
cdef int MINCHILDREN

cdef class Node:
    """ Subclasses must implement `add` and `split` methods """
    cdef public double[4] bbox
    cdef public Node parent
    # sized one larger than max to leave space for overflow before splitting
    cdef public list children
    cdef int count
    
    def __init__(self, bbox, Node parent):
        self.bbox[0] = bbox[0]
        self.bbox[1] = bbox[1]
        self.bbox[2] = bbox[2]
        self.bbox[3] = bbox[3]
        self.parent = parent
        self.count = 0

    def __len__(self):
        return self.count

cdef class LeafNode(Node):

    def add(self, MultipointBase geom):
        """ Add a geometry and return a code and a list of any new nodes
        generated by splitting. """
        cdef int retval
        cdef tuple newnodes

        retval = 0
        bb = geom.bbox
        self.bbox[0] = min(self.bbox[0], geom.bbox[0])
        self.bbox[1] = min(self.bbox[1], geom.bbox[1])
        self.bbox[2] = max(self.bbox[2], geom.bbox[2])
        self.bbox[3] = max(self.bbox[3], geom.bbox[3])
        self.children[self.count] = geom
        self.count += 1

        if self.count > MAXCHILDREN:
            retval = 1
            newnodes = self.split()
        return retval, newnodes

    def split(self, str kind):
        """ Split and return new nodes. """
        cdef int[2] seeds        

        if kind == "linear":
            seeds = linear_pick_seeds(self.children)
        elif kind == "quadratic":
            seeds = quadratic_pick_seeds(self.children)
        else:
            raise KeyError("kind must be one of 'linear', 'quadratic'")

        # add each child to a node
        cdef list children
        cdef int i, j
        cdef Node c
        cdef tuple newnodes

        children = [c for c in self.children[:self.count]]
        newnodes[0] = LeafNode(self.children[seeds[0]].bbox, self.parent)
        newnodes[1] = LeafNode(self.children[seeds[1]].bbox, self.parent)

        while self.count != 0:
            if kind == "linear":
                i, j = linear_pick_next(node0.bbox, node1.bbox, children)
            elif kind == "quadratic":
                i, j = quadratic_pick_next(node0.bbox, node1.bbox, children)
            newnodes[j].add(children[i])
            self.count -= 1

            # Ensure that each node gets a minimum number of children
            if len(newnodes[0]) + self.count == MIN_CHILDREN:
                for i in range(len(children)-1, -1, -1):
                    newnodes[0].add(children.pop(i))
            elif len(newnodes[1]) + self.count == MIN_CHILDREN:
                for i in range(len(children)-1, -1, -1):
                    newnodes[1].add(children.pop(i))

        return newnodes

cdef class NonLeafNode(Node):

    def add(self, MultipointBase geom):
        """ Add a geometry and return a code and a list of any new nodes
        generated by splitting. """
        cdef int retval, flag, i, idx
        cdef tuple newnodes, newchildnodes
        cdef Node target

        retval = 0

        # choose a child node to add to
        v = [volume_expanded(c.bbox, geom.bbox) for c in self.children]
        i = v.index(min(v))
        target = self.children[i]
        flag, newchildnodes = target.add(geom)

        self.bbox[0] = min(self.bbox[0], target.bbox[0])
        self.bbox[1] = min(self.bbox[1], target.bbox[1])
        self.bbox[2] = max(self.bbox[2], target.bbox[2])
        self.bbox[3] = max(self.bbox[3], target.bbox[3])

        if flag == 1:
            # replace
            self.children[i] = newchildnodes[0]
            self.children[self.count] = newchildnodes[1]
            self.count += 1

            if self.count > MAXCHILDREN:
                newnodes = self.split()
                retval = 1

        return retval, newnodes

    def split(self, str geom):
        """ Split and return new nodes. """
        cdef int[2] seeds

        if kind == "linear":
            seeds = linear_pick_seeds(self.children)
        elif kind == "quadratic":
            seeds = quadratic_pick_seeds(self.children)
        else:
            raise KeyError("kind must be one of 'linear', 'quadratic'")

        cdef tuple newnodes
        cdef list children

        newnodes[0] = NonLeafNode(self.children[seeds[0]].bbox, self.parent)
        newnodes[1] = NonLeafNode(self.children[seeds[1]].bbox, self.parent)

        # add each child to a node
        while self.count != 0:
            if kind == "linear":
                i, j = linear_pick_next(node0.bbox, node1.bbox, self.children)
            elif kind == "quadratic":
                i, j = quadratic_pick_next(node0.bbox, node1.bbox, self.children)

            newnodes[j].add(self.children.pop(i))
            self.count -= 1

            # Ensure that each node gets a minimum number of children
            if newnodes[0] + self.count == MIN_CHILDREN:
                for i in range(len(self.children)-1, -1, -1):
                    newnodes[0].add(self.children.pop(i))

            elif newnodes[1] + self.count == MIN_CHILDREN:
                for i in range(len(self.children)-1, -1, -1):
                    newnodes[1].add(self.children.pop(i))

        return newnodes


# raw python

def build_tree(list geoms):
    """ Construct an R-tree

    Parameters
    ----------
    geoms : list of objects implementing a `bbox` attribute
    max_children : int, optional
    """
    cdef double xmin, xmax, ymin, ymax
    cdef Node root
    cdef int flag
    cdef list newnodes 

    xmin = min(g.bbox[0] for g in geoms)
    ymin = min(g.bbox[1] for g in geoms)
    xmax = max(g.bbox[2] for g in geoms)
    ymax = max(g.bbox[3] for g in geoms)

    root = LeafNode([xmin, ymin, xmax, ymax], None)
    for geom in geoms:
        flag, newnodes = root.add(geom)
        if flag == 1:
            root = NonLeafNode(root.bbox, None)
            root.children = newnodes
    return root

def linear_pick_seeds(list geoms):
    """ Choose the seeds for splitting a list of geometries using the linear
    approximation of Guttman (1984). """
    fullbb = list(geoms[0].bbox)
    extreme_x0 = 0  # highest xmin
    extreme_x1 = 0  # lowest xmax
    extreme_y0 = 0  # highest ymin
    extreme_y1 = 0  # lowest ymax
    i = 1
    for geom in geoms[1:]:
        gbb = geom.bbox
        fullbb[0] = min(fullbb[0], gbb[0])
        fullbb[1] = min(fullbb[1], gbb[1])
        fullbb[2] = max(fullbb[2], gbb[2])
        fullbb[3] = max(fullbb[3], gbb[3])
        if geoms[extreme_x0].bbox[0] < gbb[0]:
            extreme_x0 = i
        if geoms[extreme_x1].bbox[2] > gbb[2]:
            extreme_x1 = i
        if geoms[extreme_y0].bbox[1] < gbb[1]:
            extreme_y0 = i
        if geoms[extreme_y1].bbox[3] > gbb[3]:
            extreme_y1 = i
        i += 1
    # normalize
    dx = (geoms[extreme_x0].bbox[0] - geoms[extreme_x1].bbox[2]) / (fullbb[2]-fullbb[0])
    dy = (geoms[extreme_y0].bbox[1] - geoms[extreme_y1].bbox[3]) / (fullbb[3]-fullbb[1])
    if dx > dy:
        return extreme_x0, extreme_x1
    else:
        return extreme_y0, extreme_y1

def linear_pick_next(bb0, bb1, geoms):
    """ Simply return the last geometry, as well as the index of the bbox that
    it should be placed in. """
    n = len(geoms)-1
    if volume_expanded(bb0, geoms[n].bbox) > volume_expanded(bb1, geoms[n].bbox):
        return n, 1
    else:
        return n, 0

def quadratic_pick_seeds(geoms):
    """ Choose the seeds for splitting a list of geometries using the quadratic
    approximation of Guttman (1984). """
    d = -1.0
    I, J = None, None
    for (i, j) in itertools.combinations(range(len(geoms)), 2):
        bbi = geoms[i].bbox
        bbj = geoms[j].bbox
        bbc = [min(bbi[0], bbj[0]), min(bbi[1], bbj[1]),
               max(bbi[2], bbj[2]), max(bbi[3], bbj[3])]
        d_ = area(bbc) - area(bbi) - area(bbj)
        if d_ > d:
            d = d_
            I, J = i, j
    return I, J

def quadratic_pick_next(bb0, bb1, geoms):
    """ Choose an index from *geoms* to place in one of *bb0*, *bb1*.
    Return the best entry to index next, and the bbox to place it in. """
    d0 = [volume_expanded(bb0, g.bbox) for g in geoms]
    d1 = [volume_expanded(bb1, g.bbox) for g in geoms]
    diff = [abs(a-b) for (a,b) in zip(d0, d1)]
    i = diff.index(max(diff))
    if d0[i] > d1[i]:
        return i, 1
    else:
        return i, 0

def area(bbox):
    return (bbox[2]-bbox[0])*(bbox[3]-bbox[1])

def volume_expanded(bb0, bb1):
    """ Return the volume by which the union of two bounding boxes larger than
    the bounding box of geom0. """
    union_bbox = [min(bb0[0], bb1[0]), min(bb0[1], bb1[1]),
                  max(bb0[2], bb1[2]), max(bb0[3], bb1[3])]
    return area(union_bbox)-area(bb0)

def depth(tree, d=0):
    if isinstance(tree, LeafNode):
        return d+1
    else:
        if len(tree.children) != 0:
            return max(depth(child, d+1) for child in tree.children)
        else:
            return d

# def bbox_intersection_area(bb0, bb1):
#     dx = max(min(bb0[2], bb1[2]) - max(bb0[0], bb1[0]), 0.0)
#     dy = max(min(bb0[3], bb1[3]) - max(bb0[1], bb1[1]), 0.0)
#     return dx*dy

def overlaps(bb0, bb1):
    return bbox_intersection_area(bb0, bb1) != 0.0

def search(tree, bbox):
    """ Return all geometries overlapping *bbox*. """
    if isinstance(tree, LeafNode):
        matches = []
        for geom in tree.children:
            if overlaps(geom.bbox, bbox):
                matches.append(geom)
    else:
        matches = []
        for node in tree.children:
            matches.extend(search(node, bbox))
    return matches

def contains_geom(tree, geom):
    """ Test whether tree contains *geom* """
    raise NotImplementedError()
    return

def itergeoms(tree):
    """ Iterator for all geometry ojects in tree """
    if isinstance(tree, LeafNode):
        for geom in tree.children:
            yield geom
    else:
        for child in tree.children:
            for geom in itergeoms(child):
                yield geom

def iternodes(tree):
    """ Iterator for all nodes in tree """
    for child in tree.children:
        if isinstance(child, Node):
            for _child in iternodes(child):
                yield _child
        else:
            yield child

def iterleaves(tree):
    """ Iterator for all leaf nodes in tree """
    for child in tree.children:
        if isinstance(child, NonLeafNode):
            for _child in iterleaves(child):
                yield _child
        else:
            yield child
