day17: add heatmap visualization of final path

17/src/main.rs

@@ -1,13 +1,16 @@
+use colormap::ColorMap;
+use itertools::{Itertools, MinMaxResult};
 use std::collections::hash_map::RandomState;
 use std::collections::{BinaryHeap, HashMap};
-use std::fmt::{Display, Write};
 use std::fs::File;
-use std::io::{BufRead, BufReader, Lines};
+use std::io::{BufRead, BufReader, Lines, Write};
 use std::iter::repeat;
 use std::time::Instant;
+use termcolor::{Color, ColorChoice, ColorSpec, StandardStream, WriteColor};
 
-use petgraph::algo::dijkstra;
-use petgraph::prelude::*;
+mod colormap;
+
+const COLORMAP: &ColorMap = &colormap::COLORMAP_INFERNO;
 
 // BOILERPLATE
 type InputIter = Lines<BufReader<File>>;
@@ -30,7 +33,10 @@ fn main() {
     println!("Problem 2 solution: {} [{}s]", ans2, duration.as_secs_f64());
 }
 
-// PARSE
+// DATA
+
+const UNPATH_CHAR: char = '█';
+const UNVISITED_CHAR: char = ' ';
 
 #[derive(Clone, Copy, PartialEq, Eq, Debug, Hash)]
 enum Direction {
@@ -59,6 +65,17 @@ impl Direction {
     }
 }
 
+impl From<&Direction> for char {
+    fn from(dir: &Direction) -> Self {
+        match dir {
+            Direction::Left => '←',
+            Direction::Right => '→',
+            Direction::Up => '↑',
+            Direction::Down => '↓',
+        }
+    }
+}
+
 struct CityMap {
     map: Vec<Vec<u64>>,
 }
@@ -73,6 +90,24 @@ impl CityMap {
             _ => None,
         }
     }
+    fn print(&self) -> Result<(), Box<dyn std::error::Error>> {
+        let cost_max = *self.map.iter().flat_map(|row| row.iter()).max().unwrap() as f64;
+
+        let mut so_lock = StandardStream::stdout(ColorChoice::Always);
+        for y in 0..self.map.len() {
+            for val in &self.map[y] {
+                so_lock.set_color(
+                    ColorSpec::new()
+                        .set_bg(Some(COLORMAP.apply(*val as f64 / cost_max)))
+                        .set_fg(Some(Color::Black)),
+                )?;
+                so_lock.write_fmt(format_args!("{}", val))?;
+            }
+            so_lock.reset()?;
+            writeln!(so_lock)?;
+        }
+        Ok(())
+    }
 }
 
 type Position = (usize, usize);
@@ -247,8 +282,10 @@ impl<'a> WalkCost2<'a> {
                 .collect(),
         }
     }
+    fn min_cost_at(&self, pos: Position) -> Option<&u64> {
+        self.cost_from[pos.1][pos.0].values().min()
+    }
     fn compute(&mut self, to: Position) {
-        println!("computing to: {:?}", to);
         let mut unvisited_next_moves: BinaryHeap<Move> = BinaryHeap::new();
         let valid_start_moves: Vec<Move> = Direction::all()
             .iter()
@@ -324,7 +361,7 @@ impl<'a> WalkCost2<'a> {
                                 .unwrap_or(&u64::MAX)
                     })
             }); // valid positions
-            // println!("valid moves with {:?}", cur_move);
+                // println!("valid moves with {:?}", cur_move);
             for next_move in valid_moves {
                 // println!("  {:?}", next_move);
                 unvisited_next_moves.push(next_move);
@@ -332,37 +369,63 @@ impl<'a> WalkCost2<'a> {
         }
     }
 
-
     fn shortest_path_to(&self, to: Position) -> Vec<(Position, Direction)> {
         let mut path = Vec::new();
         let mut cur_pos = to;
         // start at the end, walk backwards
         while cur_pos != self.start {
-            let (m, val) = self.cost_from[cur_pos.1][cur_pos.0].iter().min_by(|a, b| a.1.cmp(b.1)).unwrap();
+            let (m, val) = self.cost_from[cur_pos.1][cur_pos.0]
+                .iter()
+                .min_by(|a, b| a.1.cmp(b.1))
+                .unwrap();
             path.push((cur_pos, m.0));
             cur_pos = self.map.offset_pos(cur_pos, &m.0.opposite()).unwrap();
         }
         path
     }
-    fn print_shortest_path(&self, to: Position) {
+    fn print_path(&self, to: Position) -> Result<(), Box<dyn std::error::Error>> {
         let path = self.shortest_path_to(to);
         let map: HashMap<_, _, RandomState> = HashMap::from_iter(path.into_iter());
+        let cost_max_of_min = *self
+            .cost_from
+            .iter()
+            .flat_map(|row| row.iter().filter_map(|cell| cell.values().min()))
+            .max()
+            .unwrap() as f64;
+        let mut so_lock = StandardStream::stdout(ColorChoice::Always);
+
         for y in 0..self.cost_from.len() {
             for x in 0..self.map.map[y].len() {
-                if let Some(to_dir) = map.get(&(x, y)) {
-                    let c = match to_dir {
-                        Direction::Left => '<',
-                        Direction::Right => '>',
-                        Direction::Up => '^',
-                        Direction::Down => 'v',
+                let mut color = ColorSpec::new();
+                let c = if let Some(to_dir) = map.get(&(x, y)) {
+                    let normalized_cost =
+                        *self.min_cost_at((x, y)).unwrap() as f64 / cost_max_of_min;
+                    let bg_color = COLORMAP.apply(normalized_cost);
+                    let fg_color = if let Color::Rgb(r, g, b) = bg_color {
+                        Color::Rgb(255 - r, 255 - g, 255 - b)
+                    } else {
+                        Color::Black
                     };
-                    print!("{}", c);
+                    color.set_fg(Some(fg_color)).set_bg(Some(bg_color)).bold();
+                    to_dir.into()
                 } else {
-                    print!("{}", self.map.map[y][x]);
-                }
+                    if let Some(cost) = &self.min_cost_at((x, y)) {
+                        color.set_fg(Some(COLORMAP.apply(**cost as f64 / cost_max_of_min)));
+                        UNPATH_CHAR
+                    } else {
+                        color.set_fg(Some(Color::Black));
+                        UNVISITED_CHAR
+                    }
+                };
+                so_lock.set_color(&color)?;
+                let mut char_buf = [0u8; 4];
+                c.encode_utf8(&mut char_buf);
+                so_lock.write_all(&char_buf)?;
             }
-            println!();
+            so_lock.reset()?;
+            writeln!(so_lock)?;
         }
+        Ok(())
     }
 }
 
@@ -395,11 +458,14 @@ fn problem1<T: BufRead>(input: Lines<T>) -> u64 {
 // PROBLEM 2 solution
 fn problem2<T: BufRead>(input: Lines<T>) -> u64 {
     let map = CityMap::from(input);
+    // map.print().unwrap();
     let mut costs = WalkCost2::from_map(&map, (0, 0));
     costs.compute((map.map[0].len() - 1, map.map.len() - 1));
 
     // println!("{}", costs);
-    costs.print_shortest_path((costs.cost_from[0].len() - 1, costs.cost_from.len() - 1));
+    costs
+        .print_path((costs.cost_from[0].len() - 1, costs.cost_from.len() - 1))
+        .unwrap();
 
     *costs.cost_from[costs.cost_from.len() - 1][costs.cost_from[0].len() - 1]
         .values()