use kiddo::float::{distance::SquaredEuclidean, kdtree::KdTree};
use csv::ReaderBuilder;
use serde_derive::{Deserialize, Serialize};
use std::fmt;
use std::path::Path;
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct Record {
pub lat: f64,
pub lon: f64,
pub name: String,
pub admin1: String,
pub admin2: String,
pub cc: String,
}
impl Record {
pub fn as_xyz(&self) -> [f64; 3] {
degrees_lat_lng_to_unit_sphere(self.lat, self.lon)
}
}
impl fmt::Display for Record {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(
f,
"({}, {}): {}, {}, {}, {}",
self.lat, self.lon, self.name, self.admin1, self.admin2, self.cc
)
}
}
pub fn degrees_lat_lng_to_unit_sphere(lat: f64, lng: f64) -> [f64; 3] {
let lat = lat.to_radians();
let lng = lng.to_radians();
[lat.cos() * lng.cos(), lat.cos() * lng.sin(), lat.sin()]
}
#[derive(Debug, Serialize, Clone)]
pub struct SearchResult<'a> {
pub distance: f64,
pub record: &'a Record,
}
pub struct ReverseGeocoder {
locations: Vec<([f64; 2], Record)>,
tree: KdTree<f64, usize, 3, 32, u16>,
}
impl ReverseGeocoder {
pub fn new() -> ReverseGeocoder {
let mut records = Vec::new();
let cities = include_str!("../cities.csv");
let mut reader = ReaderBuilder::new()
.has_headers(true)
.from_reader(cities.as_bytes());
for record in reader.deserialize() {
let record: Record = record.unwrap();
records.push(([record.lat, record.lon], record));
}
let mut tree = KdTree::new();
records.iter().enumerate().for_each(|(idx, city)| {
tree.add(&city.1.as_xyz(), idx);
});
ReverseGeocoder {
locations: records,
tree,
}
}
pub fn from_path<P: AsRef<Path>>(file_path: P) -> Result<ReverseGeocoder, std::io::Error> {
let mut records = Vec::new();
let mut reader = ReaderBuilder::new()
.has_headers(true)
.from_path(file_path)?;
for record in reader.deserialize() {
let record: Record = record?;
records.push(([record.lat, record.lon], record));
}
if records.len() < 1 {
return Err(std::io::Error::other("Need one or more records"));
}
let mut tree = KdTree::new();
records.iter().enumerate().for_each(|(idx, city)| {
tree.add(&city.1.as_xyz(), idx);
});
Ok(ReverseGeocoder {
locations: records,
tree,
})
}
pub fn search(&self, loc: (f64, f64)) -> SearchResult {
let query = degrees_lat_lng_to_unit_sphere(loc.0, loc.1);
let nearest_neighbor = self.tree.nearest_one::<SquaredEuclidean>(&query);
let nearest = &self.locations[nearest_neighbor.item as usize];
SearchResult {
distance: nearest_neighbor.distance,
record: &nearest.1,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn it_finds_4_places() {
let geocoder = ReverseGeocoder::new();
let manhattan = geocoder.search((40.7831, -73.9712));
assert_eq!(manhattan.record.name, "Manhattan");
let slp = geocoder.search((44.962786, -93.344722));
assert_eq!(slp.record.name, "Saint Louis Park");
let mpls = geocoder.search((44.894519, -93.308702));
assert_eq!(mpls.record.name, "Richfield");
let edina = geocoder.search((44.887055, -93.334204));
assert_eq!(edina.record.name, "Edina");
}
#[test]
fn it_loads_locations_from_a_path() -> Result<(), std::io::Error> {
let geocoder = ReverseGeocoder::from_path("./cities.csv")?;
geocoder.search((45.0, 54.0));
Ok(())
}
#[test]
fn it_handles_a_nearly_blank_file() {
let geocoder = ReverseGeocoder::from_path("./nearly-blank.csv");
assert!(geocoder.is_err());
}
#[test]
fn it_handles_a_blank_file() {
let geocoder = ReverseGeocoder::from_path("./blank.csv");
assert!(geocoder.is_err());
}
#[test]
fn it_handles_an_infinite_coordinate() {
let geocoder = ReverseGeocoder::new();
geocoder.search((std::f64::INFINITY, 54.0));
}
}