use crate::bitmap::{Postings, RoarPostings};
use crate::index::InfonIndex;
use crate::tokenql::evaluate;
use serde_json::{json, Value};
use std::collections::{HashMap, HashSet};
type Ix = InfonIndex<RoarPostings>;
const CONTEXT_FACETS: &[&str] = &[
"geo", "region", "location", "loc", "place", "country", "site", "city", "date", "time", "ts",
"when", "year", "month", "qty", "unit", "value", "amount", "measure", "num", "metric", "price",
"numeric", "dur", "rel", "pol", "doctype", "kano", "sentiment", "src",
];
const MAX_STRUCT_NODES: usize = 2000;
fn facet_of(t: &str) -> &str {
t.split('/').next().unwrap_or(t)
}
fn leaf_of(t: &str) -> &str {
match t.find('/') {
Some(i) => &t[i + 1..],
None => t,
}
}
fn is_context(t: &str) -> bool {
CONTEXT_FACETS.contains(&facet_of(t))
}
fn admissible(t: &str, noise: &HashSet<String>) -> bool {
!is_context(t) && !noise.contains(t)
}
pub fn breakdown(ix: &Ix, anchor: &str, facet: &str, k: usize) -> Value {
let base = evaluate(ix, anchor);
let mut rows: Vec<(String, usize)> = Vec::new();
for t in ix.facet_members(facet) {
let n = base.and(&ix.post(t)).len();
if n > 0 {
rows.push((t.clone(), n));
}
}
rows.sort_by(|a, b| b.1.cmp(&a.1).then(a.0.cmp(&b.0)));
rows.truncate(k);
json!({
"program": "breakdown", "anchor": anchor, "facet": facet, "total": base.len(),
"partition": rows.iter().map(|(v, n)| json!({ "value": leaf_of(v), "token": v, "count": n })).collect::<Vec<_>>(),
})
}
pub fn crosstab(ix: &Ix, anchor: &str, facet_a: &str, facet_b: &str, k: usize) -> Value {
let base = evaluate(ix, anchor);
let top = |facet: &str| -> Vec<String> {
let mut v: Vec<(String, usize)> = ix
.facet_members(facet)
.into_iter()
.map(|t| (t.clone(), base.and(&ix.post(t)).len()))
.filter(|(_, n)| *n > 0)
.collect();
v.sort_by(|a, b| b.1.cmp(&a.1).then(a.0.cmp(&b.0)));
v.truncate(k);
v.into_iter().map(|(t, _)| t).collect()
};
let (a_toks, b_toks) = (top(facet_a), top(facet_b));
let matrix: Vec<Value> = a_toks
.iter()
.map(|a| {
let ba = base.and(&ix.post(a));
let cells: Vec<Value> = b_toks
.iter()
.map(|b| json!({ "col": leaf_of(b), "count": ba.and(&ix.post(b)).len() }))
.collect();
json!({ "row": leaf_of(a), "cells": cells })
})
.collect();
json!({
"program": "crosstab", "anchor": anchor, "row_facet": facet_a, "col_facet": facet_b,
"cols": b_toks.iter().map(|t| leaf_of(t)).collect::<Vec<_>>(), "matrix": matrix, "total": base.len(),
})
}
fn minmax(vals: &[f64]) -> Vec<f64> {
let (lo, hi) = vals.iter().fold((f64::INFINITY, f64::NEG_INFINITY), |(lo, hi), &v| (lo.min(v), hi.max(v)));
let rng = hi - lo;
vals.iter().map(|&v| if rng > 0.0 { (v - lo) / rng } else { 0.0 }).collect()
}
pub fn rank(ix: &Ix, facet: &str, k: usize, noise: &HashSet<String>) -> Value {
let st = Structure::build(ix, noise);
let ents: Vec<String> = ix.facet_members(facet).into_iter().cloned().collect();
if ents.is_empty() {
return json!({ "program": "rank", "facet": facet, "ranked": [] });
}
let freq: Vec<f64> = ents.iter().map(|t| ix.post_len(t) as f64).collect();
let breadth: Vec<f64> = ents.iter().map(|t| st.breadth_of(&ix.post(t)) as f64).collect();
let (nf, nb) = (minmax(&freq), minmax(&breadth));
let mut scored: Vec<(String, f64, f64, f64)> = ents
.iter()
.enumerate()
.map(|(i, t)| {
let mdus = 0.5 * nf[i] + 0.5 * nb[i];
(t.clone(), mdus, nf[i], nb[i])
})
.collect();
scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
scored.truncate(k);
json!({
"program": "rank", "facet": facet,
"ranked": scored.iter().map(|(t, m, f, b)| json!({
"token": leaf_of(t), "mdus": (m * 1000.0).round() / 1000.0,
"components": { "freq": (f * 1000.0).round() / 1000.0, "breadth": (b * 1000.0).round() / 1000.0 }
})).collect::<Vec<_>>(),
})
}
struct Structure {
names: Vec<String>,
posts: Vec<RoarPostings>,
index: HashMap<String, usize>,
forward: HashMap<u32, Vec<usize>>,
}
impl Structure {
fn build(ix: &Ix, noise: &HashSet<String>) -> Structure {
let mut nodes: Vec<(&String, usize)> = ix
.tokens()
.filter(|t| admissible(t, noise))
.map(|t| (t, ix.post_len(t)))
.filter(|(_, n)| *n > 0)
.collect();
nodes.sort_by(|a, b| b.1.cmp(&a.1).then(a.0.cmp(b.0)));
nodes.truncate(MAX_STRUCT_NODES);
let names: Vec<String> = nodes.iter().map(|(t, _)| (*t).clone()).collect();
let posts: Vec<RoarPostings> = names.iter().map(|t| ix.post(t)).collect();
let index: HashMap<String, usize> = names.iter().enumerate().map(|(i, t)| (t.clone(), i)).collect();
let mut forward: HashMap<u32, Vec<usize>> = HashMap::new();
for (i, p) in posts.iter().enumerate() {
for sid in p.to_sorted() {
forward.entry(sid).or_default().push(i);
}
}
Structure { names, posts, index, forward }
}
fn candidates(&self, post: &RoarPostings) -> HashSet<usize> {
let mut out = HashSet::new();
for sid in post.to_sorted() {
if let Some(idxs) = self.forward.get(&sid) {
out.extend(idxs.iter().copied());
}
}
out
}
fn breadth_of(&self, post: &RoarPostings) -> usize {
self.candidates(post).len()
}
fn neighbors(&self, i: usize, s: usize) -> Vec<usize> {
self.candidates(&self.posts[i])
.into_iter()
.filter(|&j| j != i && self.posts[i].and(&self.posts[j]).len() >= s)
.collect()
}
}
#[derive(Debug, Clone, serde::Serialize)]
pub struct Level {
pub s: usize,
pub primal: Graph,
pub dual: Graph,
}
#[derive(Debug, Clone, serde::Serialize)]
pub struct Graph {
pub nodes: usize,
pub edges: usize,
pub components: usize,
pub sample: Vec<(usize, usize)>,
}
fn components(n: usize, edges: &[(usize, usize)]) -> usize {
let mut parent: Vec<usize> = (0..n).collect();
fn find(p: &mut [usize], x: usize) -> usize {
let mut r = x;
while p[r] != r {
r = p[r];
}
let mut c = x;
while p[c] != r {
let next = p[c];
p[c] = r;
c = next;
}
r
}
for &(a, b) in edges {
let (ra, rb) = (find(&mut parent, a), find(&mut parent, b));
if ra != rb {
parent[ra] = rb;
}
}
let mut roots = HashSet::new();
for i in 0..n {
roots.insert(find(&mut parent, i));
}
roots.len()
}
pub fn s_filtration(ix: &Ix, max_s: usize, noise: &HashSet<String>, sample_cap: usize) -> Vec<Level> {
let tags: Vec<(&String, &RoarPostings)> = ix
.postings()
.filter(|(t, _)| !noise.contains(t.as_str()))
.collect();
let n_sit = ix.situations() as usize;
let mut tags_of: Vec<Vec<usize>> = vec![Vec::new(); n_sit];
for (ti, (_, post)) in tags.iter().enumerate() {
for sid in post.to_sorted() {
if let Some(slot) = tags_of.get_mut(sid as usize) {
slot.push(ti);
}
}
}
let overlap = |a: &[usize], b: &[usize]| -> usize {
let (mut i, mut j, mut n) = (0, 0, 0);
while i < a.len() && j < b.len() {
match a[i].cmp(&b[j]) {
std::cmp::Ordering::Equal => {
n += 1;
i += 1;
j += 1;
}
std::cmp::Ordering::Less => i += 1,
std::cmp::Ordering::Greater => j += 1,
}
}
n
};
(1..=max_s.clamp(1, 16))
.map(|s| {
let mut p_edges: Vec<(usize, usize)> = Vec::new();
for i in 0..n_sit {
for j in (i + 1)..n_sit {
if overlap(&tags_of[i], &tags_of[j]) >= s {
p_edges.push((i, j));
}
}
}
let mut d_edges: Vec<(usize, usize)> = Vec::new();
for a in 0..tags.len() {
for b in (a + 1)..tags.len() {
if tags[a].1.and(tags[b].1).len() >= s {
d_edges.push((a, b));
}
}
}
Level {
s,
primal: Graph {
nodes: n_sit,
edges: p_edges.len(),
components: components(n_sit, &p_edges),
sample: p_edges.iter().take(sample_cap).copied().collect(),
},
dual: Graph {
nodes: tags.len(),
edges: d_edges.len(),
components: components(tags.len(), &d_edges),
sample: d_edges.iter().take(sample_cap).copied().collect(),
},
}
})
.collect()
}
pub fn dual_node_names(ix: &Ix, noise: &HashSet<String>) -> Vec<String> {
ix.postings().filter(|(t, _)| !noise.contains(t.as_str())).map(|(t, _)| t.clone()).collect()
}
pub fn cooccurs(ix: &Ix, token: &str, k: usize, noise: &HashSet<String>) -> Value {
let st = Structure::build(ix, noise);
let focus = ix.post(token);
let mut scored: Vec<(usize, usize)> = st
.candidates(&focus)
.into_iter()
.filter(|&j| st.names[j] != token)
.map(|j| (j, focus.and(&st.posts[j]).len()))
.filter(|(_, n)| *n > 0)
.collect();
scored.sort_by(|a, b| b.1.cmp(&a.1).then(st.names[a.0].cmp(&st.names[b.0])));
scored.truncate(k);
json!({
"program": "structure", "op": "cooccurs", "focus": token,
"cooccurs": scored.iter().map(|(j, n)| json!({ "token": st.names[*j], "shared": n })).collect::<Vec<_>>(),
})
}
pub fn s_path(ix: &Ix, a: &str, b: &str, s: usize, noise: &HashSet<String>) -> Value {
let st = Structure::build(ix, noise);
let path = match (st.index.get(a), st.index.get(b)) {
(Some(&src), Some(&dst)) => bfs_path(&st, src, dst, s.max(1)),
_ => None,
};
json!({
"program": "structure", "op": "s_path", "a": a, "b": b, "s": s,
"path": path.map(|p| p.iter().map(|&i| st.names[i].clone()).collect::<Vec<_>>()),
})
}
fn bfs_path(st: &Structure, src: usize, dst: usize, s: usize) -> Option<Vec<usize>> {
if src == dst {
return Some(vec![src]);
}
let mut prev: HashMap<usize, Option<usize>> = HashMap::new();
prev.insert(src, None);
let mut queue = std::collections::VecDeque::from([src]);
while let Some(u) = queue.pop_front() {
for v in st.neighbors(u, s) {
if let std::collections::hash_map::Entry::Vacant(e) = prev.entry(v) {
e.insert(Some(u));
if v == dst {
let mut path = vec![dst];
let mut n = dst;
while let Some(Some(p)) = prev.get(&n) {
path.push(*p);
n = *p;
}
path.reverse();
return Some(path);
}
queue.push_back(v);
}
}
}
None
}
pub fn s_clusters(ix: &Ix, s: usize, k: usize, noise: &HashSet<String>) -> Value {
let st = Structure::build(ix, noise);
let s = s.max(1);
let mut seen = vec![false; st.names.len()];
let mut clusters: Vec<Vec<usize>> = Vec::new();
for start in 0..st.names.len() {
if seen[start] {
continue;
}
let mut comp = Vec::new();
let mut stack = vec![start];
seen[start] = true;
while let Some(u) = stack.pop() {
comp.push(u);
for v in st.neighbors(u, s) {
if !seen[v] {
seen[v] = true;
stack.push(v);
}
}
}
if comp.len() > 1 {
clusters.push(comp);
}
}
clusters.sort_by(|a, b| b.len().cmp(&a.len()));
clusters.truncate(k);
json!({
"program": "structure", "op": "s_clusters", "s": s,
"clusters": clusters.iter().map(|c| json!({
"size": c.len(),
"tokens": c.iter().take(12).map(|&i| st.names[i].clone()).collect::<Vec<_>>(),
})).collect::<Vec<_>>(),
})
}
pub fn narrow(ix: &Ix, scope: &[String], filters: &[String]) -> Value {
let mut parts: Vec<String> = Vec::new();
let mut steps: Vec<Value> = Vec::new();
let mut empty_at: Option<String> = None;
let seq: Vec<(&str, &String)> = scope
.iter()
.map(|t| ("concept", t))
.chain(filters.iter().map(|t| ("filter", t)))
.collect();
for (kind, piece) in seq {
parts.push(piece.clone());
let cur = if parts.len() == 1 { parts[0].clone() } else { format!("(and {})", parts.join(" ")) };
let n = evaluate(ix, &cur).len();
steps.push(json!({ "add": piece, "kind": kind, "expr": cur, "remaining": n }));
if n == 0 && empty_at.is_none() {
empty_at = Some(piece.clone());
}
}
let answerable = steps.last().and_then(|s| s.get("remaining")).and_then(|n| n.as_u64()).unwrap_or(0) > 0;
json!({
"program": "narrow", "steps": steps,
"verdict": if answerable { "ANSWERABLE" } else { "NOT ANSWERABLE" }, "empty_at": empty_at,
})
}