from geometry import Point

class Rect:
    def __init__(self, *args):
        if len(args) == 4:  #scalars
            self.left, self.top, self.right, self.bottom = args
        elif len(args) == 2:    #Points
            self.__init__(*(args[0].tuple() + args[1].tuple()))
        else:
            raise Exception("Need 4 scalars or 2 points, got {}".format(len(args)))

    @property
    def height(self):
        return self.bottom - self.top
    @property
    def width(self):
        return self.right - self.left

    def contains(self, p):
        return self.left <= p.x < self.right and self.top <= p.y <= self.bottom

    def overlaps(self, other):
        #returns true if line segment AB shares a common segment with line segment CD.
        def intersects(a,b,c,d):
            #returns true if x lies between a and b.
            def lies_between(x,a,b):
                return a <= x < b
            return lies_between(a, c, d) or lies_between(b, c, d) or lies_between(c, a, b) or lies_between(d, a, b)
        return intersects(self.left, self.right, other.left, other.right) and intersects(self.top, self.bottom, other.top, other.bottom)

    def copy(self, **d):
        return Rect(*[d.get(attr, getattr(self, attr)) for attr in "left top right bottom".split()])

    def corners(self):
        return (Point(self.left, self.top), Point(self.right, self.bottom))
    def tuple(self):
        return (self.left, self.top, self.right, self.bottom)

    def map(self, f):
        new_corners = [p.map(f) for p in self.corners()]
        return Rect(*new_corners)
    def pmap(self, f):
        new_corners = [f(p) for p in self.corners()]
        return Rect(*new_corners)



def bisect(rect, p):
    if rect.height > rect.width:
        return [rect.copy(bottom=p.y), rect.copy(top=p.y)]
    else:
        return [rect.copy(left=p.x), rect.copy(right=p.x)]


class KD_Tree:
    def __init__(self, rect):
        self.rect = rect
        self.p = None
        self.children = []
    def add(self, p):
        assert self.rect.contains(p)
        if self.p is None:
            self.p = p
        else:
            if not self.children:
                self.children = [KD_Tree(rect) for rect in bisect(self.rect, self.p)]
            for child in self.children:
                if child.rect.contains(p):
                    child.add(p)
                    return
            raise Exception("could not find branch for {}".format(p))
    def iter(self):
        if self.p is not None:
            yield self.p
        for child in self.children:
            for item in child.iter():
                yield item
    def iter_in_range(self, rect):
        if self.p is None:
            return
        if rect.contains(self.p):
            yield self.p
        for child in self.children:
            if rect.overlaps(child.rect):
                for item in child.iter_in_range(rect):
                    yield item


from PIL import Image, ImageDraw

def render(tree):
    #we want the resulting image to have at least one dimension that is this big
    target_size = 500

    margin = 10

    scale = target_size / float(max(tree.rect.width, tree.rect.height))

    width = int(tree.rect.width * scale) + margin*2
    height = int(tree.rect.height * scale) + margin*2
    dx = tree.rect.left
    dy = tree.rect.top
    delta = Point(dx,dy)

    img = Image.new("RGB", (width, height), "white")
    draw = ImageDraw.Draw(img)

    def logical_to_screen(p):
        return (((p - delta) * scale) + Point(margin, margin)).map(int)

    def dot(p, radius, **options):
        a = p - Point(radius, radius)
        b = p + Point(radius, radius)
        draw.chord(a.tuple() + b.tuple(), 0, 359, **options)

    def box(rect, **options):
        draw.rectangle(rect.tuple(), **options)

    def render_node(node):
        if node.p is not None:
            dot(logical_to_screen(node.p), 2, fill="black")
        for child in node.children:
            box(child.rect.pmap(logical_to_screen), outline="black")
            render_node(child)

    render_node(tree)
    radius = 1
    search_vector = Point(radius, radius)
    for p in tree.iter():
        area = Rect(p - search_vector, p + search_vector)
        for neighbor in tree.iter_in_range(area):
            draw.line(logical_to_screen(p).tuple() + logical_to_screen(neighbor).tuple(), fill="red")

    return img


def render_random_tree():
    import random
    random.seed(0)
    def rand_point(field):
        return Point(random.uniform(field.left, field.right), random.uniform(field.top, field.bottom))

    field = Rect(0,0,10,10)
    t = KD_Tree(field)

    for i in range(100):
        t.add(rand_point(field))
    render(t).save("output.png")

render_random_tree()