Geo indexes: finding what is near

// Encode interleaves the bits of a binary search on longitude (even bits) and latitude (odd bits).
// Each bit halves the range: 1 means the upper half. More bits make a smaller cell.
func Encode(lat, lng float64, bits int) uint64 {
  latLo, latHi, lngLo, lngHi := -90.0, 90.0, -180.0, 180.0
  var code uint64
  for i := range bits {
    code <<= 1
    if i%2 == 0 {
      mid := (lngLo + lngHi) / 2
      if lng >= mid {
        code |= 1
        lngLo = mid
      } else {
        lngHi = mid
      }
    } else {
      mid := (latLo + latHi) / 2
      if lat >= mid {
        code |= 1
        latLo = mid
      } else {
        latHi = mid
      }
    }
  }
  return code
}

// Decode returns the cell of a code of the given length: the box every point with that prefix
// lies in.
func Decode(code uint64, bits int) Box {
  b := Box{-90, -180, 90, 180}
  for i := range bits {
    bit := code>>(bits-1-i)&1 == 1
    if i%2 == 0 {
      mid := (b.MinLng + b.MaxLng) / 2
      if bit {
        b.MinLng = mid
      } else {
        b.MaxLng = mid
      }
    } else {
      mid := (b.MinLat + b.MaxLat) / 2
      if bit {
        b.MinLat = mid
      } else {
        b.MaxLat = mid
      }
    }
  }
  return b
}

// Neighbours returns the 8 cells around a cell, clockwise from north. A cell's neighbour can have
// a very different code: across a halving line the prefixes share nothing. So compute it from
// the box, never from the string. Past a pole there is no neighbour (ok false); past the 180°
// meridian, longitude wraps.
func Neighbours(code uint64, bits int) (cells [8]uint64, ok [8]bool) {
  b := Decode(code, bits)
  h, w := b.MaxLat-b.MinLat, b.MaxLng-b.MinLng
  lat, lng := (b.MinLat+b.MaxLat)/2, (b.MinLng+b.MaxLng)/2
  steps := [8][2]float64{{1, 0}, {1, 1}, {0, 1}, {-1, 1}, {-1, 0}, {-1, -1}, {0, -1}, {1, -1}}
  for i, s := range steps {
    nLat, nLng := lat+s[0]*h, lng+s[1]*w
    if nLat > 90 || nLat < -90 {
      continue
    }
    nLng = math.Mod(nLng+540, 360) - 180
    cells[i], ok[i] = Encode(nLat, nLng, bits), true
  }
  return cells, ok
}

// BitsFor picks the cell for a radius: the smallest cell (most bits, an even number) whose
// height and width are both at least r. Then the circle fits inside the cell and its 8
// neighbours, wherever the centre falls in its cell.
func BitsFor(r, lat float64) int {
  for bits := Bits; bits >= 2; bits -= 2 {
    if h, w := CellSize(bits, lat); h >= r && w >= r {
      return bits
    }
  }
  return 2
}

// SearchCell scans only the cell that holds the centre. It misses points across the cell border.
func (ix *Index) SearchCell(lat, lng, r float64) Search {
  bits := BitsFor(r, lat)
  c := Encode(lat, lng, bits)
  return ix.scan(lat, lng, r, bits, []uint64{c})
}

// SearchNine scans the centre cell and its 8 neighbours. The circle fits inside these 9 cells, so
// it finds every point; the distance check drops the corners.
func (ix *Index) SearchNine(lat, lng, r float64) Search {
  bits := BitsFor(r, lat)
  c := Encode(lat, lng, bits)
  cells := []uint64{c}
  ns, ok := Neighbours(c, bits)
  for i, n := range ns {
    if ok[i] {
      cells = append(cells, n)
    }
  }
  return ix.scan(lat, lng, r, bits, cells)
}

func (ix *Index) scan(lat, lng, r float64, bits int, cells []uint64) Search {
  s := Search{Bits: bits, Cells: cells}
  for _, c := range cells {
    for _, p := range ix.Cell(c, bits) {
      s.Scanned++
      if Distance(lat, lng, p.Lat, p.Lng) <= r {
        s.Found = append(s.Found, p)
      }
    }
  }
  return s
}

// Insert walks down to the leaf whose box holds p. A full leaf splits into four quarters and
// hands its points down.
func (q *Quadtree) Insert(p Point) {
  n := q.root
  for n.kids != nil {
    n = &n.kids[n.quarter(p)]
  }
  n.pts = append(n.pts, p)
  if len(n.pts) > LeafCap && n.depth < maxDepth {
    n.split()
  }
}

func (n *node) split() {
  b := n.box
  midLat, midLng := (b.MinLat+b.MaxLat)/2, (b.MinLng+b.MaxLng)/2
  n.kids = &[4]node{
    {box: Box{midLat, b.MinLng, b.MaxLat, midLng}, depth: n.depth + 1}, // north-west
    {box: Box{midLat, midLng, b.MaxLat, b.MaxLng}, depth: n.depth + 1}, // north-east
    {box: Box{b.MinLat, b.MinLng, midLat, midLng}, depth: n.depth + 1}, // south-west
    {box: Box{b.MinLat, midLng, midLat, b.MaxLng}, depth: n.depth + 1}, // south-east
  }
  pts := n.pts
  n.pts = nil
  for _, p := range pts {
    k := &n.kids[n.quarter(p)]
    k.pts = append(k.pts, p)
  }
  for i := range n.kids {
    if len(n.kids[i].pts) > LeafCap && n.kids[i].depth < maxDepth {
      n.kids[i].split()
    }
  }
}

func (n *node) quarter(p Point) int {
  q := 0
  if p.Lat < (n.box.MinLat+n.box.MaxLat)/2 {
    q = 2
  }
  if p.Lng >= (n.box.MinLng+n.box.MaxLng)/2 {
    q++
  }
  return q
}

// Within visits only the nodes whose box meets the circle's bounding box, and checks the
// distance of every point in the leaves it reaches.
func (q *Quadtree) Within(lat, lng, r float64) (found []Point, leaves []Box, scanned int) {
  bb := BoundingBox(lat, lng, r)
  var walk func(n *node)
  walk = func(n *node) {
    if !n.box.Intersects(bb) {
      return
    }
    if n.kids == nil {
      leaves = append(leaves, n.box)
      for _, p := range n.pts {
        scanned++
        if Distance(lat, lng, p.Lat, p.Lng) <= r {
          found = append(found, p)
        }
      }
      return
    }
    for i := range n.kids {
      walk(&n.kids[i])
    }
  }
  walk(q.root)
  return found, leaves, scanned
}

// Nearest returns the k points closest to (lat, lng), nearest first. One queue holds nodes, keyed
// by the distance to their box, and points, keyed by their own distance. A point that comes out
// of the queue is closer than everything still in it, so it is the next answer.
func (q *Quadtree) Nearest(lat, lng float64, k int) (found []Point, scanned int) {
  pq := &queue{{d: minDist(q.root.box, lat, lng), n: q.root}}
  for pq.Len() > 0 && len(found) < k {
    e := heap.Pop(pq).(entry)
    switch {
    case e.p != nil:
      found = append(found, *e.p)
    case e.n.kids == nil:
      for i := range e.n.pts {
        p := &e.n.pts[i]
        scanned++
        heap.Push(pq, entry{d: Distance(lat, lng, p.Lat, p.Lng), p: p})
      }
    default:
      for i := range e.n.kids {
        c := &e.n.kids[i]
        heap.Push(pq, entry{d: minDist(c.box, lat, lng), n: c})
      }
    }
  }
  return found, scanned
}