1
use std::cmp::Ordering;
2

            
3
use color_eyre::Result;
4
use petgraph::algo;
5
use petgraph::{graph::NodeIndex, Direction, Graph};
6
use svg::node::element::path::Data;
7
use svg::node::element::Element;
8
use svg::node::element::Path;
9
use svg::node::Text;
10
use svg::Document;
11
use svg::Node;
12

            
13
use crate::args::GraphArgs;
14
use crate::structs::{HapVariant, HaploNode, PhasedMatrix};
15

            
16
36
#[derive(Clone, Debug)]
17
pub struct Point {
18
    text: String,
19
    y: f32,
20
    x: f32,
21
}
22

            
23
impl std::fmt::Display for Point {
24
1
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
25
1
        let line = format!("Point {{ x: {}, y: {} }}", self.x, self.y);
26
1
        write!(f, "{line}")
27
1
    }
28
}
29

            
30
#[derive(Debug)]
31
pub struct DecayGraph<'a> {
32
    pub g: &'a Graph<HaploNode, i64>,
33
    pub vcf: &'a PhasedMatrix,
34
    pub s: GraphArgs,
35
    pub document: Document,
36
    pub tree_height: f32,
37
    // Basepairs per pixel
38
    pub scale: f32,
39
    pub nsamples: usize,
40
    pub top_padding: f32,
41
    pub side_padding: f32,
42
    pub font_size: f32,
43
    pub largest: usize,
44
    pub variables: Option<Vec<String>>,
45
}
46

            
47
impl<'a> DecayGraph<'a> {
48
42
    pub fn new(g: &'a Graph<HaploNode, i64>, vcf: &'a PhasedMatrix, s: GraphArgs) -> Self {
49
42
        let mut largest = 0;
50
1974
        for idx in g.node_indices() {
51
1974
            let node = g.node_weight(idx).unwrap();
52
1974
            // let neighbors = g.neighbors_directed(idx, Direction::Outgoing);
53
1974
            if node.pos > largest {
54
954
                largest = node.pos;
55
1020
            }
56
        }
57
42
        let nsamples = g.node_weight(NodeIndex::new(0)).unwrap().indexes.len();
58
42

            
59
42
        let top_padding = s.height as f32 * 0.02;
60
42
        let btm_padding = s.height as f32 * 0.05;
61
42
        let side_padding = 0.02;
62
42

            
63
42
        let document = Document::new()
64
42
            .set("viewBox", (0, 0, s.width, s.height))
65
42
            .set("style", format!("background-color:{}", s.background_color));
66
42

            
67
42
        Self {
68
42
            document,
69
42
            tree_height: s.height as f32 - top_padding - btm_padding,
70
42
            g,
71
42
            vcf,
72
42
            // Basepairs per pixel
73
42
            scale: largest as f32 / (s.height as f32 - top_padding - btm_padding),
74
42
            font_size: s.font_size as f32,
75
42
            s,
76
42
            nsamples,
77
42
            top_padding,
78
42
            side_padding,
79
42
            largest,
80
42
            variables: None,
81
42
        }
82
42
    }
83

            
84
684
    fn travel_majority(&self, parent_idx: NodeIndex) -> Option<NodeIndex> {
85
684
        let nodes = self
86
684
            .g
87
684
            .neighbors_directed(parent_idx, Direction::Outgoing)
88
684
            .collect::<Vec<NodeIndex>>();
89
684

            
90
684
        if nodes.len() == 2 {
91
648
            let c1 = self.g.node_weight(nodes[0]).unwrap();
92
648
            let c2 = self.g.node_weight(nodes[1]).unwrap();
93

            
94
            // Majority always to the left
95
648
            let left_index = match c1.indexes_len.cmp(&c2.indexes_len) {
96
                Ordering::Greater => nodes[0],
97
600
                Ordering::Less => nodes[1],
98
                Ordering::Equal => {
99
                    // Pick the node with gt==1
100
48
                    if let Some(last1) = c1.haplotype.iter().last() {
101
36
                        if last1.gt == 1 {
102
                            nodes[0]
103
                        } else {
104
36
                            nodes[1]
105
                        }
106
                    } else {
107
                        // This is a random choice as we cant do much better here anymore
108
12
                        nodes[0]
109
                    }
110
                },
111
            };
112

            
113
648
            self.travel_majority(left_index)
114
        } else {
115
36
            Some(parent_idx)
116
        }
117
684
    }
118

            
119
    // Finds all the nodes in the majority branch
120
24
    pub fn get_majority_nodes(&self) -> Option<Vec<NodeIndex>> {
121
24
        let start_node = NodeIndex::new(0);
122
24
        let end_node = self.travel_majority(start_node).unwrap();
123
24
        algo::all_simple_paths(&self.g, start_node, end_node, 0, None).next()
124
24
    }
125

            
126
12
    pub fn find_ancestral_haplotype(&self) -> &Vec<HapVariant> {
127
12
        let node_idx = self.travel_majority(NodeIndex::new(0)).unwrap();
128
12
        let node = self.g.node_weight(node_idx).unwrap();
129
12
        &node.haplotype
130
12
    }
131

            
132
30
    pub fn draw_graph(&mut self) -> Result<()> {
133
30
        let root_idx = NodeIndex::new(0);
134
30
        let width = (
135
30
            self.s.width as f32 * self.side_padding,
136
30
            self.s.width as f32 * (1.0 - self.side_padding),
137
30
        );
138
30

            
139
30
        // Draw once to find how much extra comes from font padding
140
30
        let mut extra_y = self.pre_draw_recursion(root_idx, width, self.top_padding, 0.0);
141
30
        extra_y *= 0.9;
142
30

            
143
30
        let n_branch_data = 3;
144
30
        extra_y += self.font_size * n_branch_data as f32;
145
30

            
146
30
        // Rescale by subtracting the pixels lost to font scaling from the tree height
147
30
        self.scale = (self.scale * self.tree_height) / (self.tree_height - extra_y);
148
30

            
149
30
        let root_node = self.g.node_weight(root_idx).unwrap();
150
30

            
151
30
        let root = Point {
152
30
            text: root_node.indexes.len().to_string(),
153
30
            x: self.s.width as f32 / 2.,
154
30
            y: self.top_padding,
155
30
        };
156
30

            
157
30
        self.document
158
30
            .append(get_text(&root, root_node, &self.s, 1.));
159
30

            
160
30
        self.recursive_draw(root_idx, root, width, self.top_padding)?;
161

            
162
30
        Ok(())
163
30
    }
164

            
165
1314
    fn recursive_draw(
166
1314
        &mut self,
167
1314
        parent_idx: NodeIndex,
168
1314
        root: Point,
169
1314
        width: (f32, f32),
170
1314
        parent_y: f32,
171
1314
    ) -> Result<()> {
172
1314
        let (start, end) = width;
173
1314
        let children: Vec<NodeIndex> = self
174
1314
            .g
175
1314
            .neighbors_directed(parent_idx, Direction::Outgoing)
176
1314
            .collect();
177
1314

            
178
1314
        if children.len() == 2 {
179
642
            let parent = self.g.node_weight(parent_idx).unwrap();
180
642
            let c1 = self.g.node_weight(children[0]).unwrap();
181
642
            let c2 = self.g.node_weight(children[1]).unwrap();
182
642

            
183
642
            let c1pos = (c1.pos - parent.pos) as f32;
184
642
            let c2pos = (c2.pos - parent.pos) as f32;
185

            
186
            // Majority always splits to the left, in equal cases gt==1 to the left
187
642
            let (left_node, right_node, left_pos, right_pos, left_index, right_index) = 
188
642
                match c1.indexes_len.cmp(&c2.indexes_len) {
189
                    Ordering::Greater => (c1, c2, c1pos, c2pos, children[0], children[1]),
190
600
                    Ordering::Less => (c2, c1, c2pos, c1pos, children[1], children[0]),
191
                    Ordering::Equal => {
192
                        // Pick the node with gt==1
193
42
                        if let Some(last1) = c1.haplotype.iter().last() {
194
36
                            if last1.gt == 1 {
195
                                (c1, c2, c1pos, c2pos, children[0], children[1])
196
                            } else {
197
36
                                (c2, c1, c2pos, c1pos, children[1], children[0])
198
                            }
199
                        } else {
200
                            // This is a random choice as we cant do much better here anymore
201
6
                                (c1, c2, c1pos, c2pos, children[0], children[1])
202
                        }
203
                    }
204
                };
205

            
206
642
            let split_pos =
207
642
                (end - start) * left_node.indexes_len as f32 / parent.indexes_len as f32;
208
642
            let split = start + split_pos;
209
642
            let x_left = start + split_pos / 2.;
210
642
            let x_right = split + (end - split) / 2.;
211
642

            
212
642
            let mut y_left = parent_y + left_pos / self.scale;
213
642
            let mut y_right = parent_y + right_pos / self.scale;
214
642

            
215
642
            if (y_left - parent_y) < self.font_size {
216
6
                y_left += self.font_size - left_pos / self.scale;
217
6
                y_right += self.font_size - left_pos / self.scale;
218
636
            }
219

            
220
642
            let left = Point {
221
642
                text: left_node.indexes_len.to_string(),
222
642
                x: x_left,
223
642
                y: y_left,
224
642
            };
225
642

            
226
642
            let right = Point {
227
642
                text: right_node.indexes_len.to_string(),
228
642
                x: x_right,
229
642
                y: y_right,
230
642
            };
231
642

            
232
642
            // let opacity = if left_node.indexes_len == 1 { 0.1 } else { 1. };
233
642

            
234
642
            self.document
235
642
                .append(get_text(&left, left_node, &self.s, 1.0));
236
642
            self.document
237
642
                .append(create_path(&root, &left, &self.s, 1.0));
238

            
239
642
            let opacity = if right_node.indexes_len == 1 { 0.1 } else { 1. };
240
642
            self.document
241
642
                .append(get_text(&right, right_node, &self.s, opacity));
242
642
            self.document
243
642
                .append(create_path(&root, &right, &self.s, opacity));
244
642

            
245
642
            // Draw data to the end of the branch
246
642
            if left_node.indexes_len == 1 {
247
36
                self.draw_branch_data(left.clone(), left_index)
248
606
            }
249

            
250
642
            self.recursive_draw(left_index, left, (start, split), y_left)?;
251
642
            self.recursive_draw(right_index, right, (split, end), y_right)?;
252
672
        }
253
1314
        Ok(())
254
1314
    }
255

            
256
36
    fn draw_branch_data(&mut self, mut point: Point, end_idx: NodeIndex) {
257
36
        let branch_end_node = self.g.node_weight(end_idx).unwrap();
258
36
        let paths =
259
36
            algo::all_simple_paths::<Vec<_>, _>(&self.g, NodeIndex::new(0), end_idx, 0, None)
260
36
                .collect::<Vec<_>>();
261
36

            
262
36
        let path = paths.first().unwrap();
263
36
        let branch_start_node = self.first_branching_node(path, end_idx).unwrap().clone();
264
36

            
265
36
        if branch_start_node.indexes_len > 5 || branch_end_node.pos == self.largest {
266
36
            let n = branch_start_node.indexes_len.to_string();
267
36

            
268
36
            // Branch end position
269
36
            let pos = branch_end_node.pos;
270
36
            point.y += self.font_size;
271
36
            point.text = format!("{:.2} mb", pos as f32 / 1_000_000.);
272
36
            self.document
273
36
                .append(branch_data(&point, &self.s, 1.0, "#ff2f8e"));
274
36

            
275
36
            // Number of branch samples
276
36
            point.y += self.font_size;
277
36
            point.text = format!("n={n}");
278
36
            self.document
279
36
                .append(branch_data(&point, &self.s, 1.0, "#66df48"));
280

            
281
36
            if let Some(variables) = &self.variables {
282
60
                for var in variables {
283
30
                    let v = self
284
30
                        .vcf
285
30
                        .get_variable_data_mean(&branch_start_node.indexes, &[var.to_owned()])
286
30
                        .unwrap();
287

            
288
                    // Age of onset mean
289
30
                    if let Some(var) = v {
290
30
                        point.y += self.font_size;
291
30
                        point.text = format!("{:.2} y", var[0]);
292
30
                        self.document
293
30
                            .append(branch_data(&point, &self.s, 1.0, "#6a77dd"));
294
30
                    }
295
                }
296
6
            }
297
        }
298
36
    }
299

            
300
36
    fn first_branching_node(&self, path: &[NodeIndex], end_idx: NodeIndex) -> Option<&HaploNode> {
301
36
        let mut prev_node = self.g.node_weight(end_idx).unwrap();
302
684
        for idx in path.iter().rev() {
303
684
            let curr_node = self.g.node_weight(*idx).unwrap();
304
684
            if curr_node.indexes_len as f32 / 2. > prev_node.indexes_len as f32 {
305
                return Some(prev_node);
306
684
            }
307
684
            prev_node = curr_node;
308
        }
309
36
        Some(self.g.node_weight(NodeIndex::new(0)).unwrap())
310
36
    }
311

            
312
594
    fn pre_draw_recursion(
313
594
        &mut self,
314
594
        node_idx: NodeIndex,
315
594
        width: (f32, f32),
316
594
        parent_y: f32,
317
594
        mut extra_y: f32,
318
594
    ) -> f32 {
319
594
        let (start, end) = width;
320
594
        let children: Vec<NodeIndex> = self
321
594
            .g
322
594
            .neighbors_directed(node_idx, Direction::Outgoing)
323
594
            .collect();
324
594

            
325
594
        if children.len() == 2 {
326
564
            let parent = self.g.node_weight(node_idx).unwrap();
327
564
            let c1 = self.g.node_weight(children[0]).unwrap();
328
564
            let c2 = self.g.node_weight(children[1]).unwrap();
329
564

            
330
564
            let c1pos = (c1.pos - parent.pos) as f32;
331
564
            let c2pos = (c2.pos - parent.pos) as f32;
332

            
333
564
            let (left_node, _, left_pos, _, left_index, _) = if c1.indexes_len > c2.indexes_len {
334
                (c1, c2, c1pos, c2pos, children[0], children[1])
335
            } else {
336
564
                (c2, c1, c2pos, c1pos, children[1], children[0])
337
            };
338

            
339
564
            let split_pos =
340
564
                (end - start) * left_node.indexes_len as f32 / parent.indexes_len as f32;
341
564

            
342
564
            let split = start + split_pos;
343
564

            
344
564
            let y_left = parent_y + left_pos / self.scale;
345
564

            
346
564
            // This means the length of the shared sequence is smaller than font size and thus
347
564
            // padding is required
348
564
            if (y_left - parent_y) < self.font_size {
349
6
                extra_y += self.font_size - left_pos / self.scale;
350
558
            }
351

            
352
            // Effectively all the font scaling comes from the majority branch,
353
            // so recursion to the right is skipped
354
564
            self.pre_draw_recursion(left_index, (start, split), y_left, extra_y)
355
        } else {
356
30
            extra_y
357
        }
358
594
    }
359
}
360

            
361
1314
pub fn get_text(p: &Point, n: &HaploNode, s: &GraphArgs, o: f32) -> Element {
362
1314
    let (color, o) = match n.decoys.is_some() && n.indexes_len == 1 {
363
        true => ("red", 1.),
364
1314
        false => (s.color.as_str(), o),
365
    };
366

            
367
1314
    let mut element = Element::new("text");
368
1314
    element.assign("x", p.x);
369
1314
    element.assign("y", p.y);
370
1314
    element.assign("opacity", o);
371
1314
    element.assign("font-weight", "bold");
372
1314
    element.assign("fill", color);
373
1314
    element.assign("font-size", format!("{}px", s.font_size));
374
1314
    element.append(Text::new(p.text.clone()));
375
1314
    element
376
1314
}
377

            
378
102
pub fn branch_data(p: &Point, s: &GraphArgs, o: f32, color: &str) -> Element {
379
102
    let mut element = Element::new("text");
380
102
    element.assign("x", p.x);
381
102
    element.assign("y", p.y);
382
102
    element.assign("opacity", o);
383
102
    element.assign("font-weight", "bold");
384
102
    element.assign("fill", color);
385
102
    element.assign("font-size", format!("{}px", s.font_size as f32 * 0.7));
386
102
    element.append(Text::new(p.text.clone()));
387
102
    element
388
102
}
389

            
390
1284
pub fn create_path(parent: &Point, child: &Point, s: &GraphArgs, o: f32) -> Path {
391
1284
    let font_size = s.font_size as f32;
392
1284
    let x_start = match parent.text.chars().count() {
393
        4 => font_size * 2.5,
394
        3 => font_size * 1.9,
395
708
        2 => font_size * 1.15,
396
576
        1 => font_size * 0.5,
397
        _ => font_size * 3.0,
398
    };
399

            
400
1284
    let shorten_path = if parent.x > child.x {
401
642
        match child.text.chars().count() {
402
            4 => font_size * 2.5,
403
            3 => font_size * 1.9,
404
318
            2 => font_size * 1.25,
405
324
            1 => font_size * 0.7,
406
            _ => font_size * 3.0,
407
        }
408
    } else {
409
642
        match child.text.chars().count() {
410
            4 => font_size * -0.74,
411
            3 => font_size * -0.54,
412
6
            2 => font_size * -0.32,
413
636
            1 => font_size * -0.10,
414
            _ => font_size * -1.0,
415
        }
416
    };
417

            
418
1284
    let data = Data::new()
419
1284
        .move_to((parent.x + x_start, parent.y + font_size * 0.18))
420
1284
        .line_to((parent.x + x_start, child.y - font_size * 0.2))
421
1284
        .line_to((child.x + shorten_path, child.y - font_size * 0.2));
422
1284

            
423
1284
    Path::new()
424
1284
        .set("fill", "none")
425
1284
        .set("stroke", s.color.as_str())
426
1284
        .set("opacity", o)
427
1284
        .set("stroke-width", s.stroke_width)
428
1284
        .set("d", data)
429
1284
}
430

            
431
#[cfg(test)]
432
#[rustfmt::skip]
433
mod tests {
434
    use super::*;
435

            
436
1
    #[test]
437
1
    fn point_display() {
438
1
        let point = Point { x: 5.0, y: 6.0, text: "".to_string() };
439
1
        assert_eq!("Point { x: 5, y: 6 }".to_string(), format!("{point}"));
440
1
    }
441
}