use crate::{
context::Context,
number::{Dimension, NumberParts},
reply::SearchReply,
};
use std::collections::BinaryHeap;
use std::{
cmp::{Ord, Ordering, PartialOrd},
collections::BTreeMap,
};
use strsim::jaro_winkler;
#[derive(PartialEq, Eq, Debug, Clone)]
pub struct SearchResult<'a> {
score: i32,
value: &'a str,
}
impl<'a> PartialOrd for SearchResult<'a> {
fn partial_cmp(&self, other: &SearchResult<'a>) -> Option<Ordering> {
self.score.partial_cmp(&other.score).map(|x| x.reverse())
}
}
impl<'a> Ord for SearchResult<'a> {
fn cmp(&self, other: &SearchResult<'a>) -> Ordering {
self.partial_cmp(other).unwrap()
}
}
pub fn search<'a>(ctx: &'a Context, query: &str, num_results: usize) -> Vec<&'a str> {
let mut results = BinaryHeap::new();
let query = query.to_lowercase();
{
let mut scan = |x: &'a str| {
let borrow = x;
let x = x.to_lowercase();
let modifier = if x == query {
4_000
} else if x.starts_with(&query) {
3_000
} else if x.ends_with(&query) {
2_000
} else if x.contains(&query) {
1_000
} else {
0_000
};
let score = jaro_winkler(&*x, &*query);
results.push(SearchResult {
score: (score * 1000.0) as i32 + modifier,
value: borrow,
});
while results.len() > num_results {
results.pop();
}
};
for k in &ctx.dimensions {
scan(&**k.id);
}
for k in ctx.units.keys() {
scan(&**k);
}
for k in ctx.quantities.values() {
scan(&**k);
}
for k in ctx.substances.keys() {
scan(&**k);
}
}
results
.into_sorted_vec()
.into_iter()
.filter_map(|x| if x.score > 800 { Some(x.value) } else { None })
.collect()
}
pub fn query(ctx: &Context, query: &str, num_results: usize) -> SearchReply {
SearchReply {
results: search(ctx, query, num_results)
.into_iter()
.map(|name| {
let parts = ctx
.lookup(name)
.map(|x| x.to_parts(ctx))
.or_else(|| {
if ctx.substances.get(name).is_some() {
Some(NumberParts {
quantity: Some("substance".to_owned()),
..Default::default()
})
} else {
None
}
})
.expect("Search returned non-existent result");
let mut raw = BTreeMap::new();
raw.insert(Dimension::new(name), 1);
NumberParts {
unit: Some(name.to_owned()),
raw_unit: Some(raw),
quantity: parts.quantity,
..Default::default()
}
})
.collect(),
}
}