use std::collections::HashMap;
use kiddo::{KdTree, SquaredEuclidean};
use nalgebra::Point2;
use serde::{Deserialize, Serialize};
use crate::lattice::{Coord, GridTransform, LatticeKind};
const GRID_TRANSFORMS_D4: [GridTransform; 8] = [
GridTransform::new(LatticeKind::Square, [[1, 0], [0, 1]], [0, 0]),
GridTransform::new(LatticeKind::Square, [[0, 1], [-1, 0]], [0, 0]),
GridTransform::new(LatticeKind::Square, [[-1, 0], [0, -1]], [0, 0]),
GridTransform::new(LatticeKind::Square, [[0, -1], [1, 0]], [0, 0]),
GridTransform::new(LatticeKind::Square, [[-1, 0], [0, 1]], [0, 0]),
GridTransform::new(LatticeKind::Square, [[1, 0], [0, -1]], [0, 0]),
GridTransform::new(LatticeKind::Square, [[0, 1], [1, 0]], [0, 0]),
GridTransform::new(LatticeKind::Square, [[0, -1], [-1, 0]], [0, 0]),
];
#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
#[non_exhaustive]
pub struct LocalMergeParams {
pub position_tol_rel: f32,
pub cell_size_ratio_tol: f32,
pub min_overlap: usize,
pub max_components: usize,
}
impl Default for LocalMergeParams {
fn default() -> Self {
Self {
position_tol_rel: 0.20,
cell_size_ratio_tol: 0.20,
min_overlap: 2,
max_components: 4,
}
}
}
#[derive(Clone, Copy, Debug)]
pub struct ComponentInput<'a> {
pub labelled: &'a HashMap<(i32, i32), usize>,
pub positions: &'a [Point2<f32>],
}
#[derive(Clone, Debug, Default)]
pub struct ComponentMergeResult {
pub components: Vec<HashMap<(i32, i32), usize>>,
pub diagnostics: ComponentMergeStats,
}
#[derive(Clone, Copy, Debug, Default)]
#[non_exhaustive]
pub struct ComponentMergeStats {
pub components_in: usize,
pub components_out: usize,
pub merges_accepted: usize,
}
fn euclidean(p: Point2<f32>, q: Point2<f32>) -> f32 {
((p.x - q.x).powi(2) + (p.y - q.y).powi(2)).sqrt()
}
fn estimate_cell_size(c: &ComponentInput<'_>) -> f32 {
let mut dists: Vec<f32> = Vec::new();
for (&(i, j), &idx) in c.labelled.iter() {
let p = c.positions[idx];
if let Some(&right) = c.labelled.get(&(i + 1, j)) {
dists.push(euclidean(p, c.positions[right]));
}
if let Some(&down) = c.labelled.get(&(i, j + 1)) {
dists.push(euclidean(p, c.positions[down]));
}
}
if dists.is_empty() {
return 0.0;
}
dists.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
dists[dists.len() / 2]
}
#[inline]
fn apply_transform(t: GridTransform, ij: (i32, i32)) -> (i32, i32) {
let v = t.apply(Coord::new(ij.0, ij.1));
(v.u, v.v)
}
fn score_alignment(
c_p: &ComponentInput<'_>,
c_q: &ComponentInput<'_>,
t: GridTransform,
delta: (i32, i32),
) -> (usize, f32) {
let mut overlap = 0usize;
let mut max_err = 0.0f32;
for (&ij_p, &idx_p) in c_p.labelled.iter() {
let ij_t = apply_transform(t, ij_p);
let ij_q = (ij_t.0 + delta.0, ij_t.1 + delta.1);
if let Some(&idx_q) = c_q.labelled.get(&ij_q) {
let err = euclidean(c_p.positions[idx_p], c_q.positions[idx_q]);
overlap += 1;
if err > max_err {
max_err = err;
}
}
}
(overlap, max_err)
}
fn find_best_alignment(
c_p: &ComponentInput<'_>,
c_q: &ComponentInput<'_>,
cell_size: f32,
params: &LocalMergeParams,
transforms: &[GridTransform],
) -> Option<(GridTransform, (i32, i32), usize)> {
let pos_tol = params.position_tol_rel * cell_size.max(1.0);
let pos_tol_sq = pos_tol * pos_tol;
let q_entries: Vec<((i32, i32), usize)> = c_q.labelled.iter().map(|(k, v)| (*k, *v)).collect();
if q_entries.is_empty() {
return None;
}
let mut tree: KdTree<f32, 2> = KdTree::new();
for (slot, (_, idx)) in q_entries.iter().enumerate() {
let pos = c_q.positions[*idx];
tree.add(&[pos.x, pos.y], slot as u64);
}
let mut hist: HashMap<(u8, i32, i32), usize> = HashMap::new();
for (&ij_p, &idx_p) in c_p.labelled.iter() {
let pos_p = c_p.positions[idx_p];
for nn in tree
.within_unsorted::<SquaredEuclidean>(&[pos_p.x, pos_p.y], pos_tol_sq)
.into_iter()
{
let slot = nn.item as usize;
let (ij_q, _idx_q) = q_entries[slot];
for (t_idx, t) in transforms.iter().enumerate() {
let tij_p = apply_transform(*t, ij_p);
let key = (t_idx as u8, ij_q.0 - tij_p.0, ij_q.1 - tij_p.1);
*hist.entry(key).or_insert(0usize) += 1;
}
}
}
let mut best: Option<(u8, (i32, i32), usize, f32)> = None;
for (&(t_idx, di, dj), &kdtree_overlap) in &hist {
if kdtree_overlap < params.min_overlap {
continue;
}
let t = transforms[t_idx as usize];
let delta = (di, dj);
let (overlap_full, max_err_full) = score_alignment(c_p, c_q, t, delta);
if overlap_full < params.min_overlap || max_err_full > pos_tol {
continue;
}
let take = match &best {
None => true,
Some((best_t_idx, best_delta, best_overlap, best_err)) => {
if overlap_full != *best_overlap {
overlap_full > *best_overlap
} else if (max_err_full - *best_err).abs() > f32::EPSILON {
max_err_full < *best_err
} else if t_idx != *best_t_idx {
t_idx < *best_t_idx
} else {
(di, dj) < *best_delta
}
}
};
if take {
best = Some((t_idx, (di, dj), overlap_full, max_err_full));
}
}
best.map(|(t_idx, d, n, _)| (transforms[t_idx as usize], d, n))
}
fn rebase(labelled: &mut HashMap<(i32, i32), usize>) {
if labelled.is_empty() {
return;
}
let min_i = labelled.keys().map(|(i, _)| *i).min().unwrap();
let min_j = labelled.keys().map(|(_, j)| *j).min().unwrap();
if min_i == 0 && min_j == 0 {
return;
}
let rebased: HashMap<(i32, i32), usize> = labelled
.drain()
.map(|((i, j), v)| ((i - min_i, j - min_j), v))
.collect();
*labelled = rebased;
}
#[cfg_attr(
feature = "tracing",
tracing::instrument(
level = "info",
skip_all,
fields(num_components = inputs.len()),
)
)]
pub fn merge_components_local(
inputs: &[ComponentInput<'_>],
params: &LocalMergeParams,
) -> ComponentMergeResult {
merge_components_local_with_transforms(inputs, params, &GRID_TRANSFORMS_D4)
}
pub fn merge_components_local_for(
inputs: &[ComponentInput<'_>],
params: &LocalMergeParams,
lattice: LatticeKind,
) -> ComponentMergeResult {
merge_components_local_with_transforms(inputs, params, lattice.symmetry_transforms())
}
fn merge_components_local_with_transforms(
inputs: &[ComponentInput<'_>],
params: &LocalMergeParams,
transforms: &[GridTransform],
) -> ComponentMergeResult {
let mut stats = ComponentMergeStats {
components_in: inputs.len(),
..Default::default()
};
if inputs.is_empty() {
return ComponentMergeResult {
components: Vec::new(),
diagnostics: stats,
};
}
let mut working: Vec<HashMap<(i32, i32), usize>> =
inputs.iter().map(|c| c.labelled.clone()).collect();
let positions_per: Vec<&[Point2<f32>]> = inputs.iter().map(|c| c.positions).collect();
let mut cell_sizes: Vec<f32> = inputs.iter().map(estimate_cell_size).collect();
let mut alive: Vec<bool> = vec![true; inputs.len()];
let mut changed = true;
while changed {
changed = false;
let mut order: Vec<usize> = (0..inputs.len()).filter(|i| alive[*i]).collect();
order.sort_by(|a, b| working[*b].len().cmp(&working[*a].len()));
'outer: for &i in &order {
for &j in &order {
if i == j || !alive[i] || !alive[j] {
continue;
}
let s_i = cell_sizes[i].max(1e-3);
let s_j = cell_sizes[j].max(1e-3);
let ratio = (s_i - s_j).abs() / s_i.max(s_j);
if ratio > params.cell_size_ratio_tol {
continue;
}
let cell_size = 0.5 * (s_i + s_j);
let c_p = ComponentInput {
labelled: &working[i],
positions: positions_per[i],
};
let c_q = ComponentInput {
labelled: &working[j],
positions: positions_per[j],
};
let Some((t, delta, _overlap)) =
find_best_alignment(&c_p, &c_q, cell_size, params, transforms)
else {
continue;
};
for (ij, idx_i) in std::mem::take(&mut working[i]) {
let tij = apply_transform(t, ij);
let key = (tij.0 + delta.0, tij.1 + delta.1);
working[j].entry(key).or_insert(idx_i);
}
alive[i] = false;
cell_sizes[j] = 0.5 * (cell_sizes[i] + cell_sizes[j]);
stats.merges_accepted += 1;
changed = true;
continue 'outer;
}
}
}
let mut out: Vec<HashMap<(i32, i32), usize>> = working
.into_iter()
.zip(alive.iter().copied())
.filter_map(|(m, a)| if a { Some(m) } else { None })
.collect();
out.sort_by_key(|m| std::cmp::Reverse(m.len()));
out.truncate(params.max_components);
for m in &mut out {
rebase(m);
}
stats.components_out = out.len();
ComponentMergeResult {
components: out,
diagnostics: stats,
}
}
#[cfg(test)]
mod tests {
use super::*;
type Labels = HashMap<(i32, i32), usize>;
type Positions = Vec<Point2<f32>>;
fn component_5x5() -> (Labels, Positions) {
let mut labelled = HashMap::new();
let mut positions = Vec::new();
for j in 0..5 {
for i in 0..5 {
let idx = positions.len();
labelled.insert((i, j), idx);
positions.push(Point2::new(i as f32 * 10.0, j as f32 * 10.0));
}
}
(labelled, positions)
}
#[test]
fn identical_components_merge_into_one() {
let (l1, p1) = component_5x5();
let (l2, p2) = component_5x5();
let inputs = vec![
ComponentInput {
labelled: &l1,
positions: &p1,
},
ComponentInput {
labelled: &l2,
positions: &p2,
},
];
let res = merge_components_local(&inputs, &LocalMergeParams::default());
assert_eq!(res.components.len(), 1);
assert_eq!(res.components[0].len(), 25);
assert_eq!(res.diagnostics.merges_accepted, 1);
}
#[test]
fn shifted_components_with_overlap_merge() {
let step = 10.0;
let mut l1 = HashMap::new();
let mut p1 = Vec::new();
for j in 0..5 {
for i in 0..3 {
let idx = p1.len();
l1.insert((i, j), idx);
p1.push(Point2::new(i as f32 * step, j as f32 * step));
}
}
let mut l2 = HashMap::new();
let mut p2 = Vec::new();
for j in 0..5 {
for i in 0..3 {
let idx = p2.len();
l2.insert((i, j), idx);
p2.push(Point2::new((i as f32 + 2.0) * step, j as f32 * step));
}
}
let inputs = vec![
ComponentInput {
labelled: &l1,
positions: &p1,
},
ComponentInput {
labelled: &l2,
positions: &p2,
},
];
let res = merge_components_local(&inputs, &LocalMergeParams::default());
assert_eq!(res.components.len(), 1);
assert_eq!(res.components[0].len(), 25);
}
#[test]
fn cell_size_mismatch_blocks_merge() {
let (l1, p1) = component_5x5();
let mut l2 = HashMap::new();
let mut p2 = Vec::new();
for j in 0..5 {
for i in 0..5 {
let idx = p2.len();
l2.insert((i, j), idx);
p2.push(Point2::new(i as f32 * 20.0, j as f32 * 20.0));
}
}
let inputs = vec![
ComponentInput {
labelled: &l1,
positions: &p1,
},
ComponentInput {
labelled: &l2,
positions: &p2,
},
];
let res = merge_components_local(&inputs, &LocalMergeParams::default());
assert_eq!(res.components.len(), 2);
assert_eq!(res.diagnostics.merges_accepted, 0);
}
#[test]
fn drifted_overlapping_corner_blocks_merge() {
let cell = 10.0_f32;
let mut l1: Labels = HashMap::new();
let mut p1: Positions = Vec::new();
for j in 0..2 {
for i in 0..2 {
let idx = p1.len();
l1.insert((i, j), idx);
p1.push(Point2::new(i as f32 * cell, j as f32 * cell));
}
}
let mut l2: Labels = HashMap::new();
let mut p2: Positions = Vec::new();
for j in 0..2 {
for i in 0..2 {
let idx = p2.len();
l2.insert((i, j), idx);
let pos = if (i, j) == (1, 1) {
Point2::new(50.0, 50.0)
} else {
Point2::new(i as f32 * cell, j as f32 * cell)
};
p2.push(pos);
}
}
let inputs = vec![
ComponentInput {
labelled: &l1,
positions: &p1,
},
ComponentInput {
labelled: &l2,
positions: &p2,
},
];
let res = merge_components_local(&inputs, &LocalMergeParams::default());
assert_eq!(
res.components.len(),
2,
"drifted corner should block the merge entirely"
);
assert_eq!(res.diagnostics.merges_accepted, 0);
}
fn hex_model(q: i32, r: i32) -> Point2<f32> {
let sqrt3_2 = 3.0_f32.sqrt() * 0.5;
Point2::new(q as f32 + 0.5 * r as f32, sqrt3_2 * r as f32)
}
fn hex_component(radius: i32, scale: f32, relabel: usize) -> (Labels, Positions) {
let t = crate::lattice::D6_TRANSFORMS[relabel];
let mut labelled = HashMap::new();
let mut positions = Vec::new();
for q in -radius..=radius {
for r in (-radius).max(-q - radius)..=radius.min(-q + radius) {
let idx = positions.len();
let m = hex_model(q, r);
positions.push(Point2::new(m.x * scale, m.y * scale));
let c = t.apply(Coord::new(q, r));
labelled.insert((c.u, c.v), idx);
}
}
(labelled, positions)
}
#[test]
fn hex_identical_components_merge_under_d6() {
let (l1, p1) = hex_component(2, 14.0, 0);
let (l2, p2) = hex_component(2, 14.0, 4); let inputs = vec![
ComponentInput {
labelled: &l1,
positions: &p1,
},
ComponentInput {
labelled: &l2,
positions: &p2,
},
];
let res =
merge_components_local_for(&inputs, &LocalMergeParams::default(), LatticeKind::Hex);
assert_eq!(
res.components.len(),
1,
"D6 merge should reunite the relabelled hex copies"
);
assert_eq!(res.components[0].len(), l1.len());
assert_eq!(res.diagnostics.merges_accepted, 1);
}
#[test]
fn hex_relabelled_copy_merges_for_every_d6_element() {
for relabel in 0..crate::lattice::D6_TRANSFORMS.len() {
let (l1, p1) = hex_component(2, 16.0, 0);
let (l2, p2) = hex_component(2, 16.0, relabel);
let inputs = vec![
ComponentInput {
labelled: &l1,
positions: &p1,
},
ComponentInput {
labelled: &l2,
positions: &p2,
},
];
let res =
merge_components_local_for(&inputs, &LocalMergeParams::default(), LatticeKind::Hex);
assert_eq!(
res.components.len(),
1,
"D6 element {relabel} failed to merge"
);
}
}
}