ck3_history_extractor/display/
graph.rs

1use rand::{rng, Rng};
2
3use plotters::{
4    coord::types::{RangedCoordf64, RangedCoordi32},
5    prelude::*,
6};
7
8// This is a cool little library that provides the TREE LAYOUT ALGORITHM, the rendering is done by plotters
9//https://github.com/zxch3n/tidy/tree/master it is sort of tiny so here github link in case it goes down
10use tidy_tree::TidyTree;
11
12use super::super::{
13    parser::GameRef,
14    structures::{Dynasty, FromGameObject, GameObjectDerived, Title},
15    types::{GameId, GameString, HashMap, Wrapper},
16};
17
18use std::{collections::BTreeMap, path::Path};
19
20/// Common graph size in pixels
21const GRAPH_SIZE: (u32, u32) = (1024, 768);
22
23/// A value indicating that the node has no parent
24const NO_PARENT: usize = usize::MAX;
25
26const GRAPH_MARGIN: u32 = 5;
27const GRAPH_LABEL_SPACE: u32 = GRAPH_MARGIN * 10;
28
29const TREE_MARGIN: u32 = 10;
30const TREE_NODE_SIZE_MULTIPLIER: f64 = 1.5;
31const PARENT_CHILD_MARGIN: f64 = 15.0;
32const PEER_MARGIN: f64 = 5.0;
33
34const TIMELINE_MARGIN: u32 = 3;
35
36/// The y label for the death graphs
37const Y_LABEL: &str = "Percentage of global deaths";
38
39/// The maximum y value for the death graphs
40const MAX_Y: f64 = 100.0;
41const MIN_Y: f64 = 0.0;
42
43/// Handles node initialization within the graph.
44/// Tree is the tree object we are adding the node to, stack is the stack we are using to traverse the tree, storage is the hashmap we are using to store the node data, fnt is the font we are using to calculate the size of the node, and parent is the parent node id.
45fn handle_node<
46    I: IntoIterator<Item = GameRef<T>>,
47    T: TreeNode<I> + GameObjectDerived + FromGameObject,
48>(
49    node: GameRef<T>,
50    tree: &mut TidyTree,
51    stack: &mut Vec<(usize, GameRef<T>)>,
52    storage: &mut HashMap<
53        usize,
54        (
55            usize,
56            GameString,
57            (f64, f64),
58            (i32, i32),
59            Option<GameString>,
60        ),
61    >,
62    parent: usize,
63    fnt: &FontDesc,
64) {
65    let obj = node.get_internal();
66    let id = obj.get_id() as usize;
67    if let Some(ch) = obj.inner() {
68        let name = ch.get_name();
69        let txt_size = fnt.box_size(&name).unwrap();
70        let node_width = txt_size.0 as f64 * TREE_NODE_SIZE_MULTIPLIER;
71        let node_height = txt_size.1 as f64 * TREE_NODE_SIZE_MULTIPLIER;
72        //we also here calculate the point where the text should be drawn while we have convenient access to both size with margin and without
73        let txt_point = (
74            -(node_width as i32 - txt_size.0 as i32),
75            -(node_height as i32 - txt_size.1 as i32),
76        );
77        //add node to tree
78        tree.add_node(id, node_width, node_height, parent);
79        stack.push((id, node.clone()));
80        //add aux data to storage because the nodes only store the id and no additional data
81        storage.insert(
82            id,
83            (
84                parent,
85                name,
86                (node_width, node_height),
87                txt_point,
88                ch.get_class(),
89            ),
90        );
91    }
92}
93
94/// A trait for objects that can be used in a tree structure
95pub trait TreeNode<I: IntoIterator>: Sized {
96    /// Returns an iterator over the children of the node
97    fn get_children(&self) -> Option<I>;
98
99    /// Returns an iterator over the parent of the node
100    fn get_parent(&self) -> Option<I>;
101
102    /// Returns the class of the node
103    fn get_class(&self) -> Option<GameString>;
104}
105
106/// Creates a graph from a given data set
107/// Assumes that the data is sorted by the x value, and that the y value is a percentage
108fn create_graph<P: AsRef<Path>, S: Into<String>>(
109    data: &BTreeMap<i16, u32>,
110    contast: &BTreeMap<i16, u32>,
111    output_path: &P,
112    ylabel: Option<S>,
113    xlabel: Option<S>,
114) {
115    let mut min_x: i16 = 0;
116    let mut max_x: i16 = 0;
117    {
118        let mut iter = data.iter();
119        if let Some((x, _)) = iter.next() {
120            min_x = *x;
121        }
122        if let Some((x, _)) = iter.next_back() {
123            max_x = *x;
124        }
125    }
126
127    let root = SVGBackend::new(output_path, GRAPH_SIZE).into_drawing_area();
128    root.fill(&WHITE).unwrap();
129    let mut chart = ChartBuilder::on(&root)
130        .margin(GRAPH_MARGIN)
131        .x_label_area_size(GRAPH_LABEL_SPACE)
132        .y_label_area_size(GRAPH_LABEL_SPACE)
133        .build_cartesian_2d((min_x as i32)..(max_x as i32), MIN_Y..MAX_Y)
134        .unwrap();
135
136    let mut mesh = chart.configure_mesh();
137
138    if let Some(xlabel) = xlabel {
139        mesh.x_desc(xlabel);
140    }
141
142    if let Some(ylabel) = ylabel {
143        mesh.y_desc(ylabel);
144    }
145
146    mesh.draw().unwrap();
147
148    chart
149        .draw_series(LineSeries::new(
150            (min_x..max_x).map(|year| {
151                (
152                    year as i32,
153                    *data.get(&year).unwrap_or(&0) as f64 / *contast.get(&year).unwrap() as f64
154                        * MAX_Y,
155                )
156            }),
157            &RED,
158        ))
159        .unwrap();
160}
161
162/// An object that can create graphs from the game state
163pub struct Grapher {
164    /// Stored graph data for all faiths, certainly less memory efficient but the speed is worth it
165    faith_graph_complete: HashMap<GameId, BTreeMap<i16, u32>>,
166    culture_graph_complete: HashMap<GameId, BTreeMap<i16, u32>>,
167    total_deaths: BTreeMap<i16, u32>,
168}
169
170impl Grapher {
171    pub fn new(
172        faith_death_data: HashMap<GameId, BTreeMap<i16, u32>>,
173        culture_death_data: HashMap<GameId, BTreeMap<i16, u32>>,
174        total_deaths: BTreeMap<i16, u32>,
175    ) -> Self {
176        Grapher {
177            faith_graph_complete: faith_death_data,
178            culture_graph_complete: culture_death_data,
179            total_deaths,
180        }
181    }
182
183    /// Creates a tree graph from a given node
184    /// The reverse parameter determines if the tree is drawn from the parent to the children or the other way around
185    pub fn create_tree_graph<
186        I: IntoIterator<Item = GameRef<T>>,
187        T: TreeNode<I> + GameObjectDerived + FromGameObject,
188        P: AsRef<Path>,
189    >(
190        &self,
191        start: GameRef<T>, // the root node that is a TreeNode
192        reverse: bool,
193        output_path: &P,
194    ) {
195        let mut tree = TidyTree::with_tidy_layout(PARENT_CHILD_MARGIN, PEER_MARGIN);
196        //tree nodes don't have any data attached to them, so we need to store the data separately
197        let mut storage = HashMap::default();
198        let fnt = ("sans-serif", 6.66 * TREE_NODE_SIZE_MULTIPLIER).into_font();
199        let mut stack = Vec::new(); //BFS stack
200        handle_node(start, &mut tree, &mut stack, &mut storage, NO_PARENT, &fnt);
201        while let Some(current) = stack.pop() {
202            if let Some(char) = current.1.get_internal().inner() {
203                let iter = if reverse {
204                    if let Some(parent) = char.get_parent() {
205                        parent
206                    } else {
207                        continue;
208                    }
209                } else {
210                    if let Some(children) = char.get_children() {
211                        children
212                    } else {
213                        continue;
214                    }
215                };
216                for el in iter {
217                    handle_node(el, &mut tree, &mut stack, &mut storage, current.0, &fnt);
218                }
219            }
220        }
221
222        tree.layout(); //this calls the layout algorithm
223
224        let root;
225
226        let mut groups: HashMap<&str, RGBColor> = HashMap::default(); //class groups
227        let mut positions = HashMap::default();
228        {
229            //convert the tree layout to a hashmap and apply a 'scale' to the drawing area
230            let layout = tree.get_pos(); //this isn't documented, but this function dumps the layout into a 3xN matrix (id, x, y)
231
232            //we need to find the area that the tree laying algorithm uses to draw the tree
233            let mut min_x = 0.0;
234            let mut max_x = 0.0;
235            let mut min_y = 0.0;
236            let mut max_y = 0.0;
237            for i in 0..layout.len() / 3 {
238                let id = layout[i * 3] as usize;
239                let x = layout[i * 3 + 1];
240                let y = layout[i * 3 + 2];
241                let (_, _, (node_width, node_height), _, class) = storage.get(&id).unwrap();
242                if let Some(class) = class {
243                    // group resolving
244                    if !groups.contains_key(class.as_ref()) {
245                        let mut rng = rng();
246                        let base: u8 = 85;
247                        let mut color = RGBColor(base, base, base);
248                        //pick a random color and make it dominant
249                        let index = rng.random_range(0..3);
250                        let add = rng.random_range(160 - base..255 - base);
251                        match index {
252                            0 => {
253                                color.0 += add;
254                            }
255                            1 => {
256                                color.1 += add;
257                            }
258                            2 => {
259                                color.2 += add;
260                            }
261                            _ => unreachable!(),
262                        }
263                        groups.insert(class.as_ref(), color);
264                    }
265                }
266                // canvas size resolving
267                positions.insert(id, (x, y));
268                let candidate_x = x - node_width;
269                if candidate_x < min_x || min_x == 0.0 {
270                    min_x = candidate_x;
271                }
272                let candidate_x = x + node_width;
273                if candidate_x > max_x {
274                    max_x = candidate_x;
275                }
276                let candidate_y = y - node_height;
277                if candidate_y < min_y {
278                    min_y = candidate_y;
279                }
280                let candidate_y = y + node_height;
281                if candidate_y > max_y {
282                    max_y = candidate_y;
283                }
284            }
285
286            let x_size = (max_x - min_x + (TREE_MARGIN as f64 * 2.0)) as u32;
287            let y_size = (max_y - min_y + (TREE_MARGIN as f64 * 2.0)) as u32;
288
289            /* NOTE on scaling
290            I did try, and I mean TRY HARD to get the scaling to work properly, but Plotters doesn't allow me to properly square rectangles.
291            Their size is in i32, which means when we try to render a tree 10k units wide the rectangle size of 0.0001 is 0.
292            This is a limitation of the library, and I can't do anything about it.
293            */
294
295            let root_raw = SVGBackend::new(output_path, (x_size, y_size)).into_drawing_area();
296
297            root_raw.fill(&WHITE).unwrap();
298
299            root = root_raw.apply_coord_spec(Cartesian2d::<RangedCoordf64, RangedCoordf64>::new(
300                min_x..max_x,
301                min_y..max_y,
302                (
303                    TREE_MARGIN as i32..(x_size - TREE_MARGIN) as i32,
304                    (y_size / 25) as i32..(y_size / 25 * 24) as i32,
305                ),
306            ));
307        }
308        //we first draw the lines. Lines go from middle points of the nodes to the middle point of the parent nodes
309        for (id, (x, y)) in &positions {
310            let (parent, _, (_, node_height), _, _) = storage.get(id).unwrap();
311            if *parent != NO_PARENT {
312                //draw the line if applicable
313                let (parent_x, parent_y) = positions.get(parent).unwrap();
314                //MAYBE improve the line laying algorithm, but it's not that important
315                root.draw(&PathElement::new(
316                    vec![
317                        (*x, *y - (node_height / 2.0)),
318                        (*parent_x, *parent_y + (node_height / 2.0)),
319                    ],
320                    Into::<ShapeStyle>::into(&BLACK).stroke_width(1),
321                ))
322                .unwrap();
323            }
324        }
325        //then we draw the nodes so that they lay on top of the lines
326        for (id, (x, y)) in &positions {
327            let (_, node_name, (node_width, node_height), txt_point, class) =
328                storage.get(id).unwrap();
329            let color = if let Some(class) = class {
330                groups.get(class.as_ref()).unwrap()
331            } else {
332                &WHITE
333            };
334            //draw the element after the line so that the line is behind the element
335            root.draw(
336                &(EmptyElement::at((*x, *y))
337                // the rectangle is defined by two points, the top left and the bottom right. We calculate the top left by subtracting half the size of the node from the center point
338                + Rectangle::new(
339                    [
340                        (-(*node_width as i32) / 2, -(*node_height as i32) / 2),
341                        (*node_width as i32 / 2, *node_height as i32 / 2)
342                    ],
343                    Into::<ShapeStyle>::into(color.mix(0.9)).filled(),
344                //we add the text to the node, the text is drawn at the point we calculated earlier
345                ) + Text::new(
346                    node_name.clone(),
347                    *txt_point,
348                    fnt.clone(),
349            )),
350            )
351            .unwrap();
352        }
353        root.present().unwrap();
354    }
355
356    /// Creates a dynasty graph, meaning the family tree graph
357    pub fn create_dynasty_graph<P: AsRef<Path>>(&self, dynasty: &Dynasty, output_path: &P) {
358        //we get the founder and use it as root
359        self.create_tree_graph(dynasty.get_founder(), false, output_path)
360    }
361
362    pub fn create_culture_graph<P: AsRef<Path>>(&self, culture_id: GameId, output_path: &P) {
363        if let Some(data) = self.culture_graph_complete.get(&culture_id) {
364            create_graph(data, &self.total_deaths, output_path, Some(Y_LABEL), None)
365        }
366    }
367
368    pub fn create_faith_graph<P: AsRef<Path>>(&self, faith_id: GameId, output_path: &P) {
369        if let Some(data) = self.faith_graph_complete.get(&faith_id) {
370            create_graph(data, &self.total_deaths, output_path, Some(Y_LABEL), None)
371        }
372    }
373}
374
375pub fn create_timeline_graph<P: AsRef<Path>>(
376    timespans: &Vec<(GameRef<Title>, Vec<(i16, i16)>)>,
377    max_date: i16,
378    output_path: P,
379) {
380    let root = SVGBackend::new(&output_path, GRAPH_SIZE).into_drawing_area();
381
382    root.fill(&WHITE).unwrap();
383
384    let t_len = timespans.len() as i32;
385    let fnt = ("sans-serif", 10.0).into_font();
386    let lifespan_y = fnt.box_size("L").unwrap().1 as i32;
387    let height = lifespan_y * t_len + TIMELINE_MARGIN as i32;
388
389    let root = root.apply_coord_spec(Cartesian2d::<RangedCoordi32, RangedCoordi32>::new(
390        0..max_date as i32,
391        -height..TIMELINE_MARGIN as i32 * 3,
392        (0..GRAPH_SIZE.0 as i32, 0..GRAPH_SIZE.1 as i32),
393    ));
394
395    root.draw(&PathElement::new(
396        [(0, 0), (max_date as i32, 0)],
397        Into::<ShapeStyle>::into(&BLACK).filled(),
398    ))
399    .unwrap();
400    const YEAR_INTERVAL: i32 = 25;
401    //draw the tick
402    for i in 0..max_date as i32 / YEAR_INTERVAL {
403        root.draw(&PathElement::new(
404            [
405                (i * YEAR_INTERVAL + 1, -height),
406                (i * YEAR_INTERVAL, TIMELINE_MARGIN as i32),
407            ],
408            Into::<ShapeStyle>::into(&BLACK).filled(),
409        ))
410        .unwrap();
411    }
412    //draw the century labels
413    for i in 1..(max_date as i32 / 100) + 1 {
414        let txt = (i * 100).to_string();
415        let txt_x = fnt.box_size(&txt).unwrap().0 as i32;
416        root.draw(&Text::new(
417            txt,
418            (i * 100 - (txt_x / 2), TIMELINE_MARGIN as i32),
419            fnt.clone(),
420        ))
421        .unwrap();
422    }
423    //draw the empire lifespans
424    for (i, (title, data)) in timespans.iter().enumerate() {
425        if let Some(title) = title.get_internal().inner() {
426            let mut txt_x = 0;
427            for (start, end) in data {
428                if *start < txt_x || txt_x == 0 {
429                    txt_x = *start;
430                }
431                let real_end;
432                if *end == 0 {
433                    real_end = max_date as i32;
434                } else {
435                    real_end = *end as i32;
436                }
437                root.draw(&Rectangle::new(
438                    [
439                        (
440                            *start as i32,
441                            -lifespan_y * i as i32 - TIMELINE_MARGIN as i32,
442                        ),
443                        (
444                            real_end,
445                            -lifespan_y * (i + 1) as i32 - TIMELINE_MARGIN as i32,
446                        ),
447                    ],
448                    Into::<ShapeStyle>::into(&GREEN).filled(),
449                ))
450                .unwrap();
451            }
452            root.draw(&Text::new(
453                title.get_name(),
454                (txt_x as i32, -lifespan_y * (i + 1) as i32),
455                fnt.clone(),
456            ))
457            .unwrap();
458        }
459    }
460    root.present().unwrap();
461}