day 17

advent_of_code::solution!(17);
use std::cmp::{max, min};

use std::collections::{HashMap, VecDeque};

#[derive(Debug, PartialEq, Eq, Hash, Copy, Clone)]
enum Dir {
    H,
    V,
}

#[derive(Debug, PartialEq, Eq, Hash, Copy, Clone)]
struct Loc {
    i: i32,
    j: i32,
    dir: Dir,
}

fn solve(input: &str, min_move: i32, max_move: i32) -> Option<u32> {
    let h = input.lines().count();
    let w = input.lines().next()?.chars().count();

    let map: Vec<Vec<u32>> = input
        .lines()
        .map(|line| line.chars().filter_map(|c| c.to_digit(10)).collect())
        .collect();

    let mut queue: VecDeque<Loc> = VecDeque::from([
        Loc { i: 0, j: 0, dir: Dir::H },
        Loc { i: 0, j: 0, dir: Dir::V },
    ]);

    let mut visited: HashMap<Loc, u32> = HashMap::from([(queue[0], 0), (queue[1], 0)]);

    while queue.len() > 0 {
        let p = queue.pop_front().unwrap();
        let loss = visited[&p];

        let search = |a, b| if a < b { (a + 1)..(b + 1) } else { b..a };

        // generate each of the possible moves from current location / direction
        match p.dir {
            Dir::H => (max(p.i - max_move, 0)..min(p.i + 1 + max_move, h as i32))
                .filter(|&i| (p.i - i).abs() >= min_move)
                .map(|i| {
                    (
                        loss + search(i, p.i)
                            .map(|ii| map[ii as usize][p.j as usize])
                            .sum::<u32>(),
                        Loc { i: i, j: p.j, dir: Dir::V },
                    )
                })
                .collect::<Vec<(u32, Loc)>>(),
            Dir::V => (max(p.j - max_move, 0)..min(p.j + 1 + max_move, w as i32))
                .filter(|&j| (p.j - j).abs() >= min_move)
                .map(|j| {
                    (
                        loss + search(j, p.j)
                            .map(|jj| map[p.i as usize][jj as usize])
                            .sum::<u32>(),
                        Loc { i: p.i, j: j, dir: Dir::H },
                    )
                })
                .collect::<Vec<(u32, Loc)>>(),
        }
        .iter()
        // store each new, and improved, move
        .for_each(|(next_loss, loc)| {
            visited
                .entry(*loc)
                .and_modify(|e| {
                    if *next_loss < *e {
                        *e = *next_loss;
                        if !queue.contains(loc) {
                            queue.push_back(*loc);
                        }
                    }
                })
                .or_insert_with(|| {
                    queue.push_back(*loc);
                    *next_loss
                });
        });
    }

    // find the best direction that ended up at the final position
    [Dir::V, Dir::H]
        .iter()
        .filter_map(|&d| {
            visited.get(&Loc {
                i: (h - 1) as i32,
                j: (w - 1) as i32,
                dir: d,
            })
        })
        .min()
        .copied()
}

pub fn part_one(input: &str) -> Option<u32> {
    solve(input, 1, 3)
}

pub fn part_two(input: &str) -> Option<u32> {
    solve(input, 4, 10)
}