tawasuyu/brahman

1//! Layout de un diagrama Sankey.
2//!
3//! Pipeline: columnas por longest-path en el DAG (back-edges descartadas)
4//! → valor de nodo = max(entrada, salida) → apilado vertical por columna
5//! con una pasada de barycenter para reducir cruces → anclas de cada
6//! banda (link) en los bordes de sus nodos.
7
8use pineal_render::{Point, Rect};
9
10/// Un nodo del Sankey.
11#[derive(Debug, Clone)]
12pub struct SankeyNode {
13    pub label: String,
14}
15
16impl SankeyNode {
17    pub fn new(label: impl Into<String>) -> Self {
18        Self { label: label.into() }
19    }
20}
21
22/// Un flujo dirigido `source → target` con un caudal `value`.
23#[derive(Debug, Clone, Copy)]
24pub struct SankeyLink {
25    pub source: usize,
26    pub target: usize,
27    pub value: f64,
28}
29
30/// Caja de un nodo ya ubicada en el lienzo.
31#[derive(Debug, Clone)]
32pub struct NodeBox {
33    pub rect: Rect,
34    pub column: usize,
35}
36
37/// Banda de un link: cuatro anclas (arriba/abajo en origen y destino).
38#[derive(Debug, Clone, Copy)]
39pub struct LinkBand {
40    pub link: usize,
41    pub src_top: Point,
42    pub src_bot: Point,
43    pub dst_top: Point,
44    pub dst_bot: Point,
45}
46
47/// Layout completo: cajas de nodos + bandas de links.
48#[derive(Debug, Clone, Default)]
49pub struct SankeyLayout {
50    pub nodes: Vec<NodeBox>,
51    pub links: Vec<LinkBand>,
52}
53
54/// Calcula el layout de un Sankey dentro de `area`.
55pub fn compute_layout(
56    nodes: &[SankeyNode],
57    links: &[SankeyLink],
58    area: Rect,
59    node_width: f32,
60    node_gap: f32,
61) -> SankeyLayout {
62    let n = nodes.len();
63    if n == 0 || area.w <= 0.0 || area.h <= 0.0 {
64        return SankeyLayout::default();
65    }
66    let valid: Vec<&SankeyLink> = links
67        .iter()
68        .filter(|l| l.source < n && l.target < n && l.source != l.target && l.value > 0.0)
69        .collect();
70
71    let columns = assign_columns(n, &valid);
72    let n_cols = columns.iter().copied().max().unwrap_or(0) + 1;
73
74    // Valor de cada nodo = max(suma entrante, suma saliente).
75    let mut in_sum = vec![0.0f64; n];
76    let mut out_sum = vec![0.0f64; n];
77    for l in &valid {
78        in_sum[l.target] += l.value;
79        out_sum[l.source] += l.value;
80    }
81    let node_value: Vec<f64> = (0..n).map(|i| in_sum[i].max(out_sum[i]).max(0.0)).collect();
82
83    // Nodos por columna.
84    let mut by_col: Vec<Vec<usize>> = vec![Vec::new(); n_cols];
85    for (i, &c) in columns.iter().enumerate() {
86        by_col[c].push(i);
87    }
88
89    // Escala vertical: la columna más cargada llena `area.h` (con gaps).
90    let max_col_value = by_col
91        .iter()
92        .map(|col| col.iter().map(|&i| node_value[i]).sum::<f64>())
93        .fold(0.0f64, f64::max)
94        .max(1e-9);
95    let max_col_count = by_col.iter().map(|c| c.len()).max().unwrap_or(1).max(1);
96    let usable_h = (area.h - node_gap * (max_col_count.saturating_sub(1)) as f32).max(1.0);
97    let v_scale = usable_h as f64 / max_col_value;
98
99    // Una pasada de barycenter para ordenar cada columna.
100    barycenter_pass(&mut by_col, &valid, &columns);
101
102    // Geometría de cada nodo.
103    let col_step = if n_cols > 1 {
104        (area.w - node_width) / (n_cols - 1) as f32
105    } else {
106        0.0
107    };
108    let mut boxes = vec![
109        NodeBox { rect: Rect::new(0.0, 0.0, 0.0, 0.0), column: 0 };
110        n
111    ];
112    for (c, col) in by_col.iter().enumerate() {
113        let mut y = area.y;
114        for &i in col {
115            let h = (node_value[i] * v_scale) as f32;
116            let x = area.x + c as f32 * col_step;
117            boxes[i] = NodeBox {
118                rect: Rect::new(x, y, node_width, h.max(1.0)),
119                column: c,
120            };
121            y += h + node_gap;
122        }
123    }
124
125    // Bandas de links: apiladas en el borde derecho del origen y el
126    // borde izquierdo del destino, en el orden de aparición.
127    let mut src_cursor = vec![0.0f32; n];
128    let mut dst_cursor = vec![0.0f32; n];
129    let mut bands = Vec::with_capacity(valid.len());
130    for (vi, l) in valid.iter().enumerate() {
131        let sb = &boxes[l.source].rect;
132        let tb = &boxes[l.target].rect;
133        let thick_s = (l.value * v_scale) as f32;
134        let s_y0 = sb.y + src_cursor[l.source];
135        let t_y0 = tb.y + dst_cursor[l.target];
136        src_cursor[l.source] += thick_s;
137        dst_cursor[l.target] += thick_s;
138        // El índice real del link en el slice original.
139        let link_idx = links
140            .iter()
141            .position(|x| {
142                x.source == l.source && x.target == l.target && x.value == l.value
143            })
144            .unwrap_or(vi);
145        bands.push(LinkBand {
146            link: link_idx,
147            src_top: Point::new(sb.right(), s_y0),
148            src_bot: Point::new(sb.right(), s_y0 + thick_s),
149            dst_top: Point::new(tb.x, t_y0),
150            dst_bot: Point::new(tb.x, t_y0 + thick_s),
151        });
152    }
153
154    SankeyLayout { nodes: boxes, links: bands }
155}
156
157/// Columna de cada nodo = longest-path desde una fuente. Las back-edges
158/// (detectadas por DFS) se descartan para romper ciclos.
159fn assign_columns(n: usize, links: &[&SankeyLink]) -> Vec<usize> {
160    let mut adj: Vec<Vec<usize>> = vec![Vec::new(); n];
161    for l in links {
162        adj[l.source].push(l.target);
163    }
164    // DFS marcando back-edges (destino en la pila actual).
165    let mut state = vec![0u8; n]; // 0=blanco 1=en-pila 2=hecho
166    let mut back: Vec<(usize, usize)> = Vec::new();
167    for s in 0..n {
168        if state[s] != 0 {
169            continue;
170        }
171        let mut stack = vec![(s, 0usize)];
172        state[s] = 1;
173        while let Some(&mut (u, ref mut i)) = stack.last_mut() {
174            if *i < adj[u].len() {
175                let v = adj[u][*i];
176                *i += 1;
177                match state[v] {
178                    0 => {
179                        state[v] = 1;
180                        stack.push((v, 0));
181                    }
182                    1 => back.push((u, v)),
183                    _ => {}
184                }
185            } else {
186                state[u] = 2;
187                stack.pop();
188            }
189        }
190    }
191    // Longest-path en el DAG (sin back-edges) vía relajación topológica.
192    let mut indeg = vec![0usize; n];
193    let mut dag: Vec<Vec<usize>> = vec![Vec::new(); n];
194    for l in links {
195        if !back.contains(&(l.source, l.target)) {
196            dag[l.source].push(l.target);
197            indeg[l.target] += 1;
198        }
199    }
200    let mut col = vec![0usize; n];
201    let mut queue: Vec<usize> = (0..n).filter(|&i| indeg[i] == 0).collect();
202    let mut head = 0;
203    while head < queue.len() {
204        let u = queue[head];
205        head += 1;
206        for &v in &dag[u] {
207            col[v] = col[v].max(col[u] + 1);
208            indeg[v] -= 1;
209            if indeg[v] == 0 {
210                queue.push(v);
211            }
212        }
213    }
214    col
215}
216
217/// Reordena los nodos de cada columna por el promedio de las posiciones
218/// de sus vecinos (barycenter heuristic), una pasada izquierda→derecha.
219fn barycenter_pass(by_col: &mut [Vec<usize>], links: &[&SankeyLink], columns: &[usize]) {
220    let n = columns.len();
221    let mut order_in_col = vec![0usize; n];
222    for col in by_col.iter() {
223        for (pos, &i) in col.iter().enumerate() {
224            order_in_col[i] = pos;
225        }
226    }
227    for c in 1..by_col.len() {
228        let bary: Vec<(usize, f64)> = by_col[c]
229            .iter()
230            .map(|&node| {
231                let mut sum = 0.0;
232                let mut cnt = 0.0;
233                for l in links {
234                    if l.target == node && columns[l.source] == c - 1 {
235                        sum += order_in_col[l.source] as f64;
236                        cnt += 1.0;
237                    }
238                }
239                (node, if cnt > 0.0 { sum / cnt } else { f64::MAX })
240            })
241            .collect();
242        let mut sorted = bary;
243        sorted.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
244        by_col[c] = sorted.into_iter().map(|(node, _)| node).collect();
245        for (pos, &i) in by_col[c].iter().enumerate() {
246            order_in_col[i] = pos;
247        }
248    }
249}
250
251#[cfg(test)]
252mod tests {
253    use super::*;
254
255    fn nodes(n: usize) -> Vec<SankeyNode> {
256        (0..n).map(|i| SankeyNode::new(format!("n{i}"))).collect()
257    }
258
259    #[test]
260    fn empty_input() {
261        let l = compute_layout(&[], &[], Rect::new(0.0, 0.0, 100.0, 100.0), 20.0, 4.0);
262        assert!(l.nodes.is_empty());
263    }
264
265    #[test]
266    fn chain_assigns_increasing_columns() {
267        // 0 → 1 → 2
268        let links = [
269            SankeyLink { source: 0, target: 1, value: 5.0 },
270            SankeyLink { source: 1, target: 2, value: 5.0 },
271        ];
272        let l = compute_layout(&nodes(3), &links, Rect::new(0.0, 0.0, 300.0, 100.0), 20.0, 4.0);
273        assert_eq!(l.nodes[0].column, 0);
274        assert_eq!(l.nodes[1].column, 1);
275        assert_eq!(l.nodes[2].column, 2);
276        assert_eq!(l.links.len(), 2);
277    }
278
279    #[test]
280    fn back_edge_does_not_loop_forever() {
281        // ciclo 0 → 1 → 0 ; debe terminar y no panickear.
282        let links = [
283            SankeyLink { source: 0, target: 1, value: 3.0 },
284            SankeyLink { source: 1, target: 0, value: 1.0 },
285        ];
286        let l = compute_layout(&nodes(2), &links, Rect::new(0.0, 0.0, 200.0, 100.0), 20.0, 4.0);
287        assert_eq!(l.nodes.len(), 2);
288    }
289
290    #[test]
291    fn node_height_proportional_to_flow() {
292        // 0 manda 10 a 1 y 1 a 2 ; nodo 0 más "grueso" que nodo 2.
293        let links = [
294            SankeyLink { source: 0, target: 1, value: 10.0 },
295            SankeyLink { source: 0, target: 2, value: 2.0 },
296        ];
297        let l = compute_layout(&nodes(3), &links, Rect::new(0.0, 0.0, 200.0, 200.0), 20.0, 4.0);
298        assert!(l.nodes[0].rect.h > l.nodes[2].rect.h);
299    }
300}