use serde::{Deserialize, Serialize};
use std::collections::HashMap;
const EARTH_RADIUS_KM: f64 = 6371.0;
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub struct GeoLocation {
pub latitude: f64,
pub longitude: f64,
}
impl GeoLocation {
#[must_use]
pub fn new(latitude: f64, longitude: f64) -> Self {
assert!(
(-90.0..=90.0).contains(&latitude),
"Latitude must be between -90 and 90 degrees"
);
assert!(
(-180.0..=180.0).contains(&longitude),
"Longitude must be between -180 and 180 degrees"
);
Self {
latitude,
longitude,
}
}
#[must_use]
#[inline]
pub fn distance_to(&self, other: &GeoLocation) -> f64 {
haversine_distance(self, other)
}
#[must_use]
#[inline]
pub fn is_within(&self, other: &GeoLocation, radius_km: f64) -> bool {
self.distance_to(other) <= radius_km
}
#[must_use]
#[inline]
pub fn bearing_to(&self, other: &GeoLocation) -> f64 {
let lat1 = self.latitude.to_radians();
let lat2 = other.latitude.to_radians();
let delta_lon = (other.longitude - self.longitude).to_radians();
let y = delta_lon.sin() * lat2.cos();
let x = lat1.cos() * lat2.sin() - lat1.sin() * lat2.cos() * delta_lon.cos();
let bearing = y.atan2(x).to_degrees();
(bearing + 360.0) % 360.0
}
}
#[must_use]
pub fn haversine_distance(loc1: &GeoLocation, loc2: &GeoLocation) -> f64 {
let lat1 = loc1.latitude.to_radians();
let lat2 = loc2.latitude.to_radians();
let delta_lat = (loc2.latitude - loc1.latitude).to_radians();
let delta_lon = (loc2.longitude - loc1.longitude).to_radians();
let a =
(delta_lat / 2.0).sin().powi(2) + lat1.cos() * lat2.cos() * (delta_lon / 2.0).sin().powi(2);
let c = 2.0 * a.sqrt().atan2((1.0 - a).sqrt());
EARTH_RADIUS_KM * c
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GeoPeer {
pub peer_id: String,
pub location: GeoLocation,
pub region: String,
pub latency_ms: f64,
pub bandwidth_mbps: f64,
}
impl GeoPeer {
#[must_use]
#[inline]
pub fn distance_to(&self, target: &GeoLocation) -> f64 {
self.location.distance_to(target)
}
#[must_use]
#[inline]
pub fn geo_score(&self, target: &GeoLocation) -> f64 {
let distance = self.distance_to(target);
let distance_score = 1.0 / (1.0 + distance / 1000.0); let latency_score = 1.0 / (1.0 + self.latency_ms / 100.0);
0.6 * distance_score + 0.4 * latency_score
}
}
#[derive(Debug, Clone)]
pub struct GeoConfig {
pub preferred_radius_km: f64,
pub max_distance_km: f64,
pub enable_region_grouping: bool,
pub min_region_diversity: usize,
pub distance_weight: f64,
}
impl Default for GeoConfig {
fn default() -> Self {
Self {
preferred_radius_km: 500.0, max_distance_km: 10000.0, enable_region_grouping: true,
min_region_diversity: 2, distance_weight: 0.6, }
}
}
pub struct GeoSelector {
config: GeoConfig,
peers: HashMap<String, GeoPeer>,
regions: HashMap<String, Vec<String>>,
}
impl GeoSelector {
#[must_use]
pub fn new(config: GeoConfig) -> Self {
Self {
config,
peers: HashMap::new(),
regions: HashMap::new(),
}
}
pub fn add_peer(&mut self, peer: GeoPeer) {
self.regions
.entry(peer.region.clone())
.or_default()
.push(peer.peer_id.clone());
self.peers.insert(peer.peer_id.clone(), peer);
}
pub fn remove_peer(&mut self, peer_id: &str) -> Option<GeoPeer> {
if let Some(peer) = self.peers.remove(peer_id) {
if let Some(region_peers) = self.regions.get_mut(&peer.region) {
region_peers.retain(|id| id != peer_id);
if region_peers.is_empty() {
self.regions.remove(&peer.region);
}
}
Some(peer)
} else {
None
}
}
#[must_use]
#[inline]
pub fn find_nearest(&self, target: &GeoLocation, n: usize) -> Vec<GeoPeer> {
let mut peers_with_distance: Vec<(GeoPeer, f64)> = self
.peers
.values()
.map(|peer| {
let distance = peer.distance_to(target);
(peer.clone(), distance)
})
.collect();
peers_with_distance.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
peers_with_distance
.into_iter()
.take(n)
.map(|(peer, _)| peer)
.collect()
}
#[must_use]
#[inline]
pub fn find_within_radius(&self, target: &GeoLocation, radius_km: f64) -> Vec<GeoPeer> {
self.peers
.values()
.filter(|peer| peer.distance_to(target) <= radius_km)
.cloned()
.collect()
}
#[must_use]
#[inline]
pub fn select_best(&self, target: &GeoLocation, n: usize) -> Vec<GeoPeer> {
let mut peers_with_score: Vec<(GeoPeer, f64)> = self
.peers
.values()
.map(|peer| {
let score = peer.geo_score(target);
(peer.clone(), score)
})
.collect();
peers_with_score.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
peers_with_score
.into_iter()
.take(n)
.map(|(peer, _)| peer)
.collect()
}
#[must_use]
#[inline]
pub fn select_diverse(&self, target: &GeoLocation, n: usize) -> Vec<GeoPeer> {
if !self.config.enable_region_grouping {
return self.select_best(target, n);
}
let mut selected = Vec::new();
let mut used_regions = std::collections::HashSet::new();
let mut peers_with_score: Vec<(GeoPeer, f64)> = self
.peers
.values()
.map(|peer| {
let score = peer.geo_score(target);
(peer.clone(), score)
})
.collect();
peers_with_score.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
for (peer, _) in &peers_with_score {
if !used_regions.contains(&peer.region) {
selected.push(peer.clone());
used_regions.insert(peer.region.clone());
if selected.len() >= n {
return selected;
}
}
}
for (peer, _) in peers_with_score {
if selected.iter().any(|p| p.peer_id == peer.peer_id) {
continue;
}
selected.push(peer);
if selected.len() >= n {
break;
}
}
selected
}
#[must_use]
#[inline]
pub fn get_peers_by_region(&self, region: &str) -> Vec<GeoPeer> {
if let Some(peer_ids) = self.regions.get(region) {
peer_ids
.iter()
.filter_map(|id| self.peers.get(id))
.cloned()
.collect()
} else {
Vec::new()
}
}
#[must_use]
#[inline]
pub fn get_regions(&self) -> Vec<String> {
self.regions.keys().cloned().collect()
}
#[must_use]
pub fn get_geo_stats(&self) -> GeoStats {
let mut region_counts = HashMap::new();
for (region, peers) in &self.regions {
region_counts.insert(region.clone(), peers.len());
}
let total_peers = self.peers.len();
let total_regions = self.regions.len();
GeoStats {
total_peers,
total_regions,
peers_per_region: region_counts,
avg_peers_per_region: if total_regions > 0 {
total_peers as f64 / total_regions as f64
} else {
0.0
},
}
}
#[must_use]
#[inline]
pub fn peer_count(&self) -> usize {
self.peers.len()
}
#[must_use]
#[inline]
pub fn get_peer(&self, peer_id: &str) -> Option<&GeoPeer> {
self.peers.get(peer_id)
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GeoStats {
pub total_peers: usize,
pub total_regions: usize,
pub peers_per_region: HashMap<String, usize>,
pub avg_peers_per_region: f64,
}
#[must_use]
pub fn midpoint(loc1: &GeoLocation, loc2: &GeoLocation) -> GeoLocation {
let lat1 = loc1.latitude.to_radians();
let lon1 = loc1.longitude.to_radians();
let lat2 = loc2.latitude.to_radians();
let lon2 = loc2.longitude.to_radians();
let bx = lat2.cos() * (lon2 - lon1).cos();
let by = lat2.cos() * (lon2 - lon1).sin();
let lat3 = (lat1.sin() + lat2.sin()).atan2(((lat1.cos() + bx).powi(2) + by.powi(2)).sqrt());
let lon3 = lon1 + by.atan2(lat1.cos() + bx);
GeoLocation {
latitude: lat3.to_degrees(),
longitude: lon3.to_degrees(),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_geo_location_creation() {
let loc = GeoLocation::new(37.7749, -122.4194);
assert_eq!(loc.latitude, 37.7749);
assert_eq!(loc.longitude, -122.4194);
}
#[test]
#[should_panic]
fn test_invalid_latitude() {
let _ = GeoLocation::new(91.0, 0.0);
}
#[test]
fn test_haversine_distance() {
let sf = GeoLocation::new(37.7749, -122.4194);
let ny = GeoLocation::new(40.7128, -74.0060);
let distance = sf.distance_to(&ny);
assert!((distance - 4130.0).abs() < 50.0);
}
#[test]
fn test_same_location_distance() {
let loc = GeoLocation::new(0.0, 0.0);
assert_eq!(loc.distance_to(&loc), 0.0);
}
#[test]
fn test_is_within_radius() {
let loc1 = GeoLocation::new(37.7749, -122.4194);
let loc2 = GeoLocation::new(37.3382, -121.8863);
assert!(loc1.is_within(&loc2, 100.0));
assert!(!loc1.is_within(&loc2, 50.0));
}
#[test]
fn test_geo_selector() {
let config = GeoConfig::default();
let mut selector = GeoSelector::new(config);
let peer1 = GeoPeer {
peer_id: "peer1".to_string(),
location: GeoLocation::new(37.7749, -122.4194),
region: "us-west".to_string(),
latency_ms: 50.0,
bandwidth_mbps: 100.0,
};
let peer2 = GeoPeer {
peer_id: "peer2".to_string(),
location: GeoLocation::new(40.7128, -74.0060),
region: "us-east".to_string(),
latency_ms: 120.0,
bandwidth_mbps: 100.0,
};
selector.add_peer(peer1);
selector.add_peer(peer2);
assert_eq!(selector.peer_count(), 2);
let target = GeoLocation::new(37.3382, -121.8863);
let nearest = selector.find_nearest(&target, 1);
assert_eq!(nearest.len(), 1);
assert_eq!(nearest[0].peer_id, "peer1");
}
#[test]
fn test_region_grouping() {
let config = GeoConfig::default();
let mut selector = GeoSelector::new(config);
selector.add_peer(GeoPeer {
peer_id: "peer1".to_string(),
location: GeoLocation::new(37.7749, -122.4194),
region: "us-west".to_string(),
latency_ms: 50.0,
bandwidth_mbps: 100.0,
});
let region_peers = selector.get_peers_by_region("us-west");
assert_eq!(region_peers.len(), 1);
let regions = selector.get_regions();
assert!(regions.contains(&"us-west".to_string()));
}
#[test]
fn test_midpoint() {
let loc1 = GeoLocation::new(0.0, 0.0);
let loc2 = GeoLocation::new(0.0, 10.0);
let mid = midpoint(&loc1, &loc2);
assert!((mid.latitude - 0.0).abs() < 0.01);
assert!((mid.longitude - 5.0).abs() < 0.01);
}
#[test]
fn test_bearing() {
let loc1 = GeoLocation::new(0.0, 0.0);
let loc2 = GeoLocation::new(1.0, 0.0);
let bearing = loc1.bearing_to(&loc2);
assert!((bearing - 0.0).abs() < 1.0 || (bearing - 360.0).abs() < 1.0);
}
}