advent_of_code_2024/src/util/maps.rs

use std::fmt::Display;

use itertools::Itertools;

#[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)]
pub enum Direction {
    North,
    South,
    East,
    West,
}

impl Direction {
    pub fn all_variants() -> Vec<Self> {
        vec![Self::North, Self::East, Self::West, Self::South]
    }

    pub fn opposite(self) -> Self {
        match self {
            Self::North => Self::South,
            Self::East => Self::West,
            Self::West => Self::East,
            Self::South => Self::North,
        }
    }

    pub fn is_parallel(self, other: Self) -> bool {
        self == other || self == other.opposite()
    }

    pub fn clockwise(self) -> Self {
        match self {
            Self::North => Self::East,
            Self::East => Self::South,
            Self::South => Self::West,
            Self::West => Self::North,
        }
    }

    pub fn counterclockwise(self) -> Self {
        match self {
            Self::North => Self::West,
            Self::West => Self::South,
            Self::South => Self::East,
            Self::East => Self::North,
        }
    }
}

#[derive(Clone)]
pub struct Map<T>
where
    T: Copy,
{
    pub map: Vec<Vec<T>>,
}

impl Map<char> {
    pub fn from_string(string: &str) -> Self {
        Map::from_2d_vec(
            string
                .lines()
                .filter(|line| !line.is_empty())
                .map(|line| line.chars().collect())
                .collect(),
        )
    }
}

impl<T> Map<T>
where
    T: Copy,
{
    pub fn from_2d_vec(map: Vec<Vec<T>>) -> Self {
        Self { map }
    }

    pub fn from_dimensions(width: usize, height: usize, default: T) -> Self {
        Self {
            map: vec![vec![default; width]; height],
        }
    }

    pub fn travel(&self, x: usize, y: usize, direction: Direction) -> Option<(usize, usize)> {
        match direction {
            Direction::North => Some((x, y.checked_sub(1)?)),
            Direction::South => {
                if y + 1 < self.height() {
                    Some((x, y + 1))
                } else {
                    None
                }
            }
            Direction::West => Some((x.checked_sub(1)?, y)),
            Direction::East => {
                if x + 1 < self.width() {
                    Some((x + 1, y))
                } else {
                    None
                }
            }
        }
    }

    pub fn get(&self, x: usize, y: usize) -> Option<T> {
        Some(*self.map.get(y)?.get(x)?)
    }

    pub fn travel_get(&self, x: usize, y: usize, direction: Direction) -> Option<T> {
        let (nx, ny) = self.travel(x, y, direction)?;
        self.get(nx, ny)
    }

    pub fn travel_and_get(
        &self,
        x: usize,
        y: usize,
        direction: Direction,
    ) -> Option<(usize, usize, T)> {
        let (nx, ny) = self.travel(x, y, direction)?;
        Some((nx, ny, self.get(nx, ny)?))
    }

    pub fn set(&mut self, x: usize, y: usize, value: T) {
        self.map[y][x] = value
    }

    pub fn height(&self) -> usize {
        self.map.len()
    }

    pub fn width(&self) -> usize {
        self.map[0].len()
    }

    pub fn in_bounds(&self, x: usize, y: usize) -> bool {
        x < self.width() && y < self.height()
    }

    pub fn iter(&self) -> impl Iterator<Item = &T> {
        self.map.iter().flat_map(|row| row.iter())
    }

    pub fn iter_rows(&self) -> impl Iterator<Item = &Vec<T>> {
        self.map.iter()
    }

    pub fn coordinates(&self) -> impl Iterator<Item = (usize, usize)> {
        (0..self.width()).cartesian_product(0..self.height())
    }

    pub fn enumerate(&self) -> impl Iterator<Item = (usize, usize, &T)> {
        self.coordinates().map(|(x, y)| (x, y, &self.map[y][x]))
    }

    pub fn map<U: Copy, F: Fn(&T) -> U>(&self, f: F) -> Map<U> {
        let v = self.map.iter().map(|row| row.iter().map(|e| f(e)).collect()).collect();
        Map::<U>::from_2d_vec(v)
    }
}

impl<T> Display for Map<T>
where
    T: Copy + Display,
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let mut final_string = "".to_string();
        for row in self.map.iter() {
            let s = row.iter().fold("".to_string(), |mut acc, e| {
                acc.push_str(&format!("{}", e));
                acc
            });
            final_string.push_str("\n");
            final_string.push_str(&s);
        }
        f.write_str(&final_string)
    }
}