ck3_history_extractor/game_data/
map.rs

1use std::{borrow::Borrow, error, io, num::ParseIntError, ops::Deref, path::Path, thread};
2
3use csv::ReaderBuilder;
4
5use derive_more::{Display, From};
6use image::{
7    imageops::{crop_imm, resize, FilterType},
8    ImageReader, Rgb, RgbImage,
9};
10
11use plotters::{
12    backend::BitMapBackend,
13    drawing::IntoDrawingArea,
14    element::Text,
15    prelude::{EmptyElement, Rectangle},
16    style::{Color, IntoFont, RGBAColor, ShapeStyle, BLACK},
17};
18use serde::Serialize;
19
20use base64::prelude::*;
21
22use super::super::types::{GameId, GameString, HashMap};
23
24// color stuff
25
26/// The color of the text drawn on the map
27const TEXT_COLOR: RGBAColor = RGBAColor(0, 0, 0, 0.5);
28/// The color of the water on the map
29const WATER_COLOR: Rgb<u8> = Rgb([20, 150, 255]);
30/// The color of the land on the map
31const LAND_COLOR: Rgb<u8> = Rgb([255, 255, 255]);
32/// The color of the null pixels on the map
33const NULL_COLOR: Rgb<u8> = Rgb([0, 0, 0]);
34
35// map image stuff
36
37/// The scale factor for the input map image
38const SCALE: u32 = 4;
39
40#[derive(Debug, From, Display)]
41pub enum MapError {
42    IoError(io::Error),
43    ImageError(image::ImageError),
44    DefinitionError(csv::Error),
45    ParsingError(ParseIntError),
46}
47
48impl error::Error for MapError {
49    fn source(&self) -> Option<&(dyn error::Error + 'static)> {
50        match self {
51            MapError::IoError(e) => Some(e),
52            MapError::ImageError(e) => Some(e),
53            MapError::DefinitionError(e) => Some(e),
54            MapError::ParsingError(e) => Some(e),
55        }
56    }
57}
58
59/// Returns a vector of bytes from a png file encoded with rgb8, meaning each pixel is represented by 3 bytes
60fn read_png_bytes<P: AsRef<Path>>(path: P) -> Result<RgbImage, MapError> {
61    Ok(ImageReader::open(path)?.decode()?.to_rgb8())
62}
63
64pub trait MapImage {
65    fn draw_text<T: Borrow<str>>(&mut self, text: T);
66
67    fn draw_legend<C: Into<Rgb<u8>>, I: IntoIterator<Item = (String, C)>>(&mut self, legend: I);
68
69    fn save_in_thread<P: AsRef<Path>>(self, path: P);
70}
71
72impl MapImage for RgbImage {
73    /// Draws given text on an image buffer, the text is placed at the bottom left corner and is 5% of the height of the image
74    fn draw_text<T: Borrow<str>>(&mut self, text: T) {
75        let dimensions = self.dimensions();
76        let text_height = dimensions.1 / 20;
77        let back = BitMapBackend::with_buffer(self, dimensions).into_drawing_area();
78        let style = ("sans-serif", text_height).into_font().color(&TEXT_COLOR);
79        back.draw(&Text::new(
80            text,
81            (10, dimensions.1 as i32 - text_height as i32),
82            style,
83        ))
84        .unwrap();
85        back.present().unwrap();
86    }
87
88    /// Draws a legend on the given image buffer, the legend is placed at the bottom right corner and consists of a series of colored rectangles with text labels
89    fn draw_legend<C: Into<Rgb<u8>>, I: IntoIterator<Item = (String, C)>>(&mut self, legend: I) {
90        let dimensions = self.dimensions();
91        let text_height = (dimensions.1 / 30) as i32;
92        let back = BitMapBackend::with_buffer(self, dimensions).into_drawing_area();
93        let style = ("sans-serif", text_height).into_font();
94        let mut x = (dimensions.0 / 50) as i32;
95        for (label, color) in legend {
96            let text_size = style.box_size(&label).unwrap();
97            let margin = text_height / 3;
98            let color = color.into();
99            back.draw(
100                &(EmptyElement::at((x, dimensions.1 as i32 - (text_height * 2)))
101                    + Rectangle::new(
102                        [(0, 0), (text_height, text_height)],
103                        ShapeStyle {
104                            color: RGBAColor(color[0], color[1], color[2], 1.0),
105                            filled: true,
106                            stroke_width: 1,
107                        },
108                    )
109                    + Rectangle::new(
110                        [(0, 0), (text_height, text_height)],
111                        ShapeStyle {
112                            color: BLACK.to_rgba(),
113                            filled: false,
114                            stroke_width: 1,
115                        },
116                    )
117                    + Text::new(
118                        label,
119                        (text_height + margin, (text_height - text_size.1 as i32)),
120                        style.clone(),
121                    )),
122            )
123            .unwrap();
124            x += text_height + text_size.0 as i32 + (margin * 2);
125        }
126        back.present().unwrap();
127    }
128
129    fn save_in_thread<P: AsRef<Path>>(self, path: P) {
130        let path = path.as_ref().to_owned();
131        thread::spawn(move || {
132            self.save(path).unwrap();
133        });
134    }
135}
136
137fn serialize_into_b64<S: serde::Serializer>(
138    image: &RgbImage,
139    serializer: S,
140) -> Result<S::Ok, S::Error> {
141    serializer.serialize_str(&BASE64_STANDARD.encode(image.as_raw()))
142}
143
144fn serialize_title_color_map<S: serde::Serializer>(
145    title_color_map: &HashMap<GameString, Rgb<u8>>,
146    serializer: S,
147) -> Result<S::Ok, S::Error> {
148    let mut map = HashMap::new();
149    for (k, v) in title_color_map.iter() {
150        map.insert(k.clone(), v.0);
151    }
152    serializer.serialize_some(&map)
153}
154
155/// A struct representing a game map, from which we can create [Map] instances
156#[derive(Serialize)]
157pub struct GameMap {
158    height: u32,
159    width: u32,
160    #[serde(serialize_with = "serialize_into_b64")]
161    province_map: RgbImage,
162    #[serde(serialize_with = "serialize_title_color_map")]
163    title_color_map: HashMap<GameString, Rgb<u8>>,
164}
165
166impl GameMap {
167    /// Creates a new GameMap from a province map and a definition.csv file located inside the provided path.
168    /// The function expects the path to be a valid CK3 game directory.
169    pub fn new<P: AsRef<Path>>(
170        provinces_path: P,
171        rivers_path: P,
172        definition_path: P,
173        province_barony_map: &HashMap<GameId, GameString>,
174    ) -> Result<Self, MapError> {
175        let mut provinces = read_png_bytes(provinces_path)?;
176        let river = read_png_bytes(rivers_path)?;
177        //apply river bytes as a mask to the provinces bytes so that non white pixels in rivers are black
178        for (p_b, r_b) in provinces.pixels_mut().zip(river.pixels()) {
179            if r_b[0] != 255 || r_b[1] != 255 || r_b[2] != 255 {
180                *p_b = NULL_COLOR;
181            }
182        }
183        let (width, height) = provinces.dimensions();
184        // we need to find a bounding box for the terrain
185        let mut max_x = 0;
186        let mut min_x = width;
187        let mut max_y = 0;
188        let mut min_y = height;
189        for x in 0..width {
190            for y in 0..height {
191                if provinces.get_pixel(x, y) != &NULL_COLOR {
192                    if x > max_x {
193                        max_x = x;
194                    }
195                    if x < min_x {
196                        min_x = x;
197                    }
198                    if y > max_y {
199                        max_y = y;
200                    }
201                    if y < min_y {
202                        min_y = y;
203                    }
204                }
205            }
206        }
207        let width = max_x - min_x;
208        let height = max_y - min_y;
209        let cropped = crop_imm(&provinces, min_x, min_y, width, height);
210
211        //scale the image down to 1/4 of the size
212        provinces = resize(
213            cropped.deref(),
214            (width / SCALE) as u32,
215            (height / SCALE) as u32,
216            FilterType::Nearest,
217        );
218        //provinces.save("test.png").unwrap();
219        //ok so now we have a province map with each land province being a set color and we now just need to read definition.csv
220        let mut key_colors: HashMap<GameString, Rgb<u8>> = HashMap::default();
221        let mut rdr = ReaderBuilder::new()
222            .comment(Some(b'#'))
223            .flexible(true)
224            .delimiter(b';')
225            .from_path(definition_path)?;
226        for record in rdr.records() {
227            let record = match record {
228                Ok(record) => record,
229                Err(_) => continue,
230            };
231            let id = match record[0].parse::<GameId>() {
232                Ok(id) => id,
233                Err(_) => continue,
234            };
235            let r = record[1].parse::<u8>()?;
236            let g = record[2].parse::<u8>()?;
237            let b = record[3].parse::<u8>()?;
238            if let Some(barony) = province_barony_map.get(&id) {
239                key_colors.insert(barony.clone(), Rgb([r, g, b]));
240            }
241        }
242        Ok(GameMap {
243            height: height,
244            width: width,
245            province_map: provinces,
246            title_color_map: key_colors,
247        })
248    }
249}
250
251pub trait MapGenerator {
252    /// Creates a new instance of map, with pixels colored in accordance with assoc function. key_list acts as a whitelist of keys to use assoc on. If it's None then assoc is applied to all keys.
253    fn create_map<C: Into<Rgb<u8>> + Clone, F: Fn(&str) -> C, I: IntoIterator<Item = GameString>>(
254        &self,
255        assoc: F,
256        key_list: Option<I>,
257    ) -> RgbImage;
258
259    /// Creates a new instance of map, with all pixels corresponding to keys in the key_list colored same as the target_color
260    fn create_map_flat<C: Into<Rgb<u8>> + Clone, I: IntoIterator<Item = GameString>>(
261        &self,
262        key_list: I,
263        target_color: C,
264    ) -> RgbImage {
265        self.create_map(|_| target_color.clone(), Some(key_list))
266    }
267}
268
269impl MapGenerator for GameMap {
270    // this place is very hot according to the perf profiler
271    fn create_map<
272        C: Into<Rgb<u8>> + Clone,
273        F: Fn(&str) -> C,
274        I: IntoIterator<Item = GameString>,
275    >(
276        &self,
277        assoc: F,
278        key_list: Option<I>,
279    ) -> RgbImage {
280        let mut colors = HashMap::default();
281        if let Some(key_list) = key_list {
282            for k in key_list {
283                if let Some(color) = self.title_color_map.get(k.as_ref()) {
284                    colors.insert(color, assoc(k.as_ref()));
285                } else {
286                    // huh? this is weird
287                }
288            }
289        } else {
290            for (k, v) in self.title_color_map.iter() {
291                colors.insert(v, assoc(k));
292            }
293        }
294        let mut new_map = self.province_map.clone();
295        for pixel in new_map.pixels_mut() {
296            if pixel == &NULL_COLOR {
297                *pixel = WATER_COLOR;
298            } else {
299                if let Some(color) = colors.get(pixel) {
300                    *pixel = (*color).clone().into();
301                } else {
302                    *pixel = LAND_COLOR;
303                }
304            }
305        }
306        return new_map;
307    }
308}