Geohelpers.swift

import Foundation
import MapKit

extension MKMapPoint {
    static var nyc: MKMapPoint {
        return MKMapPoint(.nyc)
    }

}
extension MKMapRect {
    var northWestPoint: MKMapPoint {
        MKMapPoint(x: minX, y: maxY)
    }

    var northEastPoint: MKMapPoint {
        MKMapPoint(x: maxX, y: maxY)
    }

    var southWestPoint: MKMapPoint {
        MKMapPoint(x: minX, y: minY)
    }

    var southEastPoint: MKMapPoint {
        MKMapPoint(x: maxX, y: minY)
    }

    var center: MKMapPoint {
        get {
            MKMapPoint(x: midX, y: midY)
        }
        set {
            self.origin = MKMapPoint(x: newValue.x - width/2, y: newValue.y - height/2)
        }
    }

    var quadrants: (northWest: MKMapRect, northEast: MKMapRect, southWest: MKMapRect, southEast: MKMapRect) {
        var north: MKMapRect = .null
        var south: MKMapRect = .null

        MKMapRectDivide(self, &north, &south, height/2, .minYEdge)

        var northWest: MKMapRect = .null
        var northEast: MKMapRect = .null
        var southWest: MKMapRect = .null
        var southEast: MKMapRect = .null

        MKMapRectDivide(north, &northWest, &northEast, width/2, .minXEdge)
        MKMapRectDivide(south, &southWest, &southEast, width/2, .minXEdge)

        return (northWest, northEast, southWest, southEast)
    }
}

QuadTree.swift

protocol MapPointHaving {
    var mapPoint: MKMapPoint { get }
}

final class QuadTreeNode<Element: MapPointHaving & Identifiable> {

    enum Children {
        case undivided([Element])
        case divided(northWest: QuadTreeNode, northEast: QuadTreeNode, southWest: QuadTreeNode, southEast: QuadTreeNode)

        var areDivided: Bool {
            if case .divided = self {
                return true
            }
            return false
        }

        mutating func divideIfNeeded(rect: MKMapRect) -> (northWest: QuadTreeNode, northEast: QuadTreeNode, southWest: QuadTreeNode, southEast: QuadTreeNode) {
            if case .undivided = self {
                let quadrants = rect.quadrants
                let northWest = QuadTreeNode(rect: quadrants.northWest)
                let northEast = QuadTreeNode(rect: quadrants.northEast)
                let southWest = QuadTreeNode(rect: quadrants.southWest)
                let southEast = QuadTreeNode(rect: quadrants.southEast)

                self = .divided(northWest: northWest, northEast: northEast, southWest: southWest, southEast: southEast)
            }
            return assumingDivided()
        }

        func ifDivided() -> (northWest: QuadTreeNode, northEast: QuadTreeNode, southWest: QuadTreeNode, southEast: QuadTreeNode)? {
            if case let .divided(northWest, northEast, southWest, southEast) = self {
                return (northWest, northEast, southWest, southEast)
            } else {
                return nil
            }
        }


        func assumingDivided() -> (northWest: QuadTreeNode, northEast: QuadTreeNode, southWest: QuadTreeNode, southEast: QuadTreeNode) {
            if let divided = ifDivided() {
                return divided
            } else {
                fatalError("Tried to assume that the QuadTree's Children was divided, but it was not")
            }
        }
    }

    var children: Children = .undivided([])

    var rect: MKMapRect

    let bucketSize: Int = 8

    init(rect: MKMapRect) {
        self.rect = rect
    }

    @discardableResult
    func insert(_ value: Element) -> Bool {
        guard self.rect.contains(value.mapPoint) else {
            return false
        }

        switch self.children {
        case let .undivided(points):
            if points.count < bucketSize {
                self.children = .undivided(points + [value])
                return true
            }

            let quadrants = self.children.divideIfNeeded(rect: rect)

            for point in points {
                quadrants.northWest.insert(point)
                quadrants.northEast.insert(point)
                quadrants.southWest.insert(point)
                quadrants.southEast.insert(point)
            }

            // could this potentially result in one child having more than bucketSize children?
            return quadrants.northWest.insert(value) || quadrants.northEast.insert(value) || quadrants.southWest.insert(value) || quadrants.southEast.insert(value)

        case let .divided(northWest, northEast, southWest, southEast):
             return northWest.insert(value) || northEast.insert(value) || southWest.insert(value) || southEast.insert(value)
        }
    }

    func elements(in rect: MKMapRect) -> [Element] {
        guard self.rect.intersects(rect) else {
            return []
        }
        switch self.children {
        case let .undivided(points):
            return points.filter({ rect.contains($0.mapPoint) })
        case let .divided(northWest, northEast, southWest, southEast):
            return northWest.elements(in: rect)
            + northEast.elements(in: rect)
            + southWest.elements(in: rect)
            + southEast.elements(in: rect)
        }
    }

    func closestPlaces(to point: MKMapPoint, maxCount: Int = 20, closerThan distanceInMeters: Double) -> [RelativePlace<Element>] {

        if !rect.contains(point)
            && rect.northWestPoint.distance(to: point) > distanceInMeters
            && rect.northEastPoint.distance(to: point) > distanceInMeters
            && rect.southWestPoint.distance(to: point) > distanceInMeters
            && rect.southEastPoint.distance(to: point) > distanceInMeters {
            return []
        }

        let maxDistance = distanceInMeters

        // do we need this??
        //        var maxDistance: Double {
        //            if points.count >= bucketSize, let farthest = points.max(by: { $0.distance < $1.distance }) {
        //                return farthest.distance
        //            } else {
        //                return .greatestFiniteMagnitude
        //            }
        //        }

        switch self.children {
        case let .undivided(points):
            return points
                .map({ RelativePlace(place: $0, relativeTo: point) })
                .smallest(maxCount)
        case let .divided(nw, ne, sw, se):
            return (nw.closestPlaces(to: point, maxCount: maxCount, closerThan: maxDistance)
                + ne.closestPlaces(to: point, maxCount: maxCount, closerThan: maxDistance)
                + sw.closestPlaces(to: point, maxCount: maxCount, closerThan: maxDistance)
                + se.closestPlaces(to: point, maxCount: maxCount, closerThan: maxDistance))
                .smallest(maxCount)
        }
    }
}

// extension QuadTreeNode where Element == PlaceAnnotation {
//     func update(_ place: Place) {
//         //check which quadrant it should go in
//         switch self.children {
//         case var .undivided(points):
//             if let index = points.firstIndex(where: { place.id == $0.id }) {
//                 let annotation = PlaceAnnotation(place: place)
//                 points[index] = annotation
//                 self.children = .undivided(points)
//             }
//         case let .divided(nw, ne, sw, se):
//             ne.update(place)
//             nw.update(place)
//             se.update(place)
//             sw.update(place)
//         }
//     }
//
//     func delete(_ place: Place) {
//         switch self.children {
//         case var .undivided(points):
//             if let index = points.firstIndex(where: { place.id == $0.id }) {
//                 points.remove(at: index)
//                 self.children = .undivided(points)
//             }
//         case let .divided(nw, ne, sw, se):
//             ne.delete(place)
//             nw.delete(place)
//             se.delete(place)
//             sw.delete(place)
//         }
//     }
// }
//

extension QuadTreeNode: Collection {

    /// Traversal order - self, northwest, northeast, southwest, southeast
    struct Index: Comparable {

        enum Branch: Comparable {
            case nw, ne, sw, se

            var sortOrder: Int {
                switch self {
                case .nw:
                    return 0
                case .ne:
                    return 1
                case .sw:
                    return 2
                case .se:
                    return 3
                }
            }

            static func < (lhs: Branch, rhs: Branch) -> Bool {
                lhs.sortOrder < rhs.sortOrder
            }
        }

        struct NodeIndex {
            // should nodes be weak?
            let branchPath: [(parent: QuadTreeNode, direction: Branch)]

            let currentNode: QuadTreeNode

            func nextNonEmptyUndividedNode() -> NodeIndex? {
                return sequence(first: self, next: { $0.nextUndividedNode() })
                    .dropFirst()
                    .first(where: {
                        if case let .undivided(points) = $0.currentNode.children, !points.isEmpty {
                            return true
                        } else {
                            return false
                        }
                    })
            }

            func nextUndividedNode() -> NodeIndex? {
                switch branchPath.last {
                case let (parent, .nw)?:
                    var path = Array(branchPath.dropLast()) + [(parent, .ne)]
                    var node = parent.child(in: .ne)!
                    while let quadrants = node.children.ifDivided() {
                        path.append((node, .nw))
                        node = quadrants.northWest
                    }
                    return NodeIndex(branchPath: path, currentNode: node)
                case let (parent, .ne)?:
                    var path = Array(branchPath.dropLast()) + [(parent, .sw)]
                    var node = parent.child(in: .sw)!
                    while let quadrants = node.children.ifDivided() {
                        path.append((node, .nw))
                        node = quadrants.northWest
                    }
                    return NodeIndex(branchPath: path, currentNode: node)
                case let (parent, .sw)?:
                    var path = Array(branchPath.dropLast()) + [(parent, .se)]
                    var node = parent.child(in: .se)!
                    while let quadrants = node.children.ifDivided() {
                        path.append((node, .nw))
                        node = quadrants.northWest
                    }
                    return NodeIndex(branchPath: path, currentNode: node)
                case let (parent, .se)?:
                    return NodeIndex(branchPath: Array(branchPath.dropLast()), currentNode: parent)
                        .nextUndividedNode()
                case nil:
                    return nil
                }
            }
        }

        var nodeIndex: NodeIndex

        let arrayIndex: Int

        init(branchPath: [(parent: QuadTreeNode, direction: Branch)], currentNode: QuadTreeNode, arrayIndex: Int) {
            self.nodeIndex = NodeIndex(branchPath: branchPath, currentNode: currentNode)
            self.arrayIndex = arrayIndex
        }

        var branchPath: [(parent: QuadTreeNode, direction: Branch)] {
            nodeIndex.branchPath
        }

        var currentNode: QuadTreeNode {
            nodeIndex.currentNode
        }

        static func < (lhs: Index, rhs: Index) -> Bool {
            let lhsDirections = lhs.branchPath.lazy.map({ $0.direction })
            let rhsDirections = rhs.branchPath.lazy.map({ $0.direction })
            if lhsDirections.elementsEqual(rhsDirections) {
                return lhs.arrayIndex < rhs.arrayIndex
            } else {
                return lhsDirections.lexicographicallyPrecedes(rhsDirections)
            }
        }

        static func == (lhs: Index, rhs: Index) -> Bool {
            let lhsDirections = lhs.branchPath.lazy.map({ $0.direction })
            let rhsDirections = rhs.branchPath.lazy.map({ $0.direction })
            return lhsDirections.elementsEqual(rhsDirections) && lhs.arrayIndex == rhs.arrayIndex
        }
    }

    var startIndex: Index {
        var current = self
        var path: [(QuadTreeNode, Index.Branch)] = []
        while case let .divided(nw, _, _, _) = current.children {
            path.append((current, .nw))
            current = nw
        }

        guard let index = Index.NodeIndex(branchPath: path, currentNode: current).nextNonEmptyUndividedNode() else { return endIndex }
        return Index(branchPath: index.branchPath, currentNode: index.currentNode, arrayIndex: 0)
    }

    var endIndex: Index {
        var current = self
        var path: [(QuadTreeNode, Index.Branch)] = []
        while case let .divided(_, _, _, se) = current.children {
            path.append((current, .se))
            current = se
        }
        return Index(branchPath: path, currentNode: current, arrayIndex: Int.max)
    }

    func index(after i: Index) -> Index {
        guard case let .undivided(points) = i.currentNode.children else {
            fatalError("Every index should point to a leaf node")
        }

        if i.arrayIndex < points.count - 1 {
            return Index(branchPath: i.branchPath, currentNode: i.currentNode, arrayIndex: i.arrayIndex + 1)
        }

        guard let j = i.nodeIndex.nextNonEmptyUndividedNode() else {
            return endIndex
        }
        return Index(branchPath: j.branchPath, currentNode: j.currentNode, arrayIndex: 0)
    }

    subscript(position: Index) -> Element {
        if case let .undivided(points) = position.currentNode.children {
            return points[position.arrayIndex]
        } else {
            fatalError("out of bounds!")
        }
    }

    func node<C: Collection>(for path: C) -> QuadTreeNode? where C.Element == Index.Branch {
        guard let firstBranch = path.first else { return self }

        //invalid path, maybe should crash by changing ? to !
        return self.child(in: firstBranch)?.node(for: path.dropFirst())
    }

    func child(in direction: Index.Branch) -> QuadTreeNode? {
        if let quadrants = self.children.ifDivided() {
            switch direction {
            case .nw:
                return quadrants.northWest
            case .ne:
                return quadrants.northEast
            case .sw:
                return quadrants.southWest
            case .se:
                return quadrants.southEast
            }
        }
        return nil
    }

}

extension QuadTreeNode {
    func asArray() -> [Element] {
        switch self.children {
        case let .undivided(points):
            return points
        case let .divided(nw, ne, sw, se):
            return nw.asArray() + ne.asArray() + sw.asArray() + se.asArray()
        }
    }
}

extension Sequence {
    func smallest(_ n: Int, usingTest areInIncreasingOrder: (Element, Element) -> Bool) -> [Element] {
        var result = self.prefix(n).sorted(by: areInIncreasingOrder)

        for e in self.dropFirst(n) {
            if let last = result.last, areInIncreasingOrder(last, e) { continue }
            if let insertionIndex = result.firstIndex(where: { areInIncreasingOrder(e, $0) }) {
                result.insert(e, at: insertionIndex)
                result.removeLast()
            }
        }

        return result
    }
}

extension Sequence where Element: Comparable {
    func smallest(_ n: Int) -> [Element] {
        self.smallest(n, usingTest: <)
    }

}