2024 day 21

use memoize::memoize;

fn numpad(key: char) -> (i32, i32) {
    match key {
        '7' => (0, 0),
        '8' => (0, 1),
        '9' => (0, 2),
        '4' => (1, 0),
        '5' => (1, 1),
        '6' => (1, 2),
        '1' => (2, 0),
        '2' => (2, 1),
        '3' => (2, 2),
        '0' => (3, 1),
        'A' => (3, 2),
        _ => panic!(),
    }
}

fn arrowpad(key: char) -> (i32, i32) {
    match key {
        '^' => (0, 1),
        'A' => (0, 2),
        '<' => (1, 0),
        'v' => (1, 1),
        '>' => (1, 2),
        _ => panic!(),
    }
}

#[memoize]
fn do_arrows(i: i32, j: i32, steps: usize, h_first: bool) -> usize {
    let (ii, jj) = (i.unsigned_abs() as usize, j.unsigned_abs() as usize);
    let mut chunk = vec![if i > 0 { '^' } else { 'v' }; ii];
    chunk.extend(vec![if j > 0 { '<' } else { '>' }; jj]);

    if h_first {
        chunk.reverse();
    }

    chunk.push('A');

    if steps == 0 {
        chunk.len()
    } else {
        let mut loc = arrowpad('A');

        chunk
            .into_iter()
            .map(|c| {
                let n = arrowpad(c);
                let p = loc;
                loc = n;
                let d = (p.0 - n.0, p.1 - n.1);
                if d.0 == 0 || d.1 == 0 {
                    // straight line, search only once, order is irrelevant
                    do_arrows(d.0, d.1, steps - 1, false)
                } else if n == (1, 0) && p.0 == 0 {
                    // must search down first
                    do_arrows(d.0, d.1, steps - 1, false)
                } else if p == (1, 0) && n.0 == 0 {
                    // must search horiz first
                    do_arrows(d.0, d.1, steps - 1, true)
                } else {
                    // can search in either order
                    std::cmp::min(
                        do_arrows(d.0, d.1, steps - 1, false),
                        do_arrows(d.0, d.1, steps - 1, true),
                    )
                }
            })
            .sum()
    }
}

fn enter_sequence(sequence: &str, steps: usize) -> usize {
    let mut loc = numpad('A');

    sequence
        .chars()
        .map(|c| {
            // either move horizontally or vertically first
            // in some cases only one will be valid...
            let n = numpad(c);
            let p = loc;
            let d = (loc.0 - n.0, loc.1 - n.1);
            loc = n;
            if p.0 == 3 && n.1 == 0 {
                // must move up first
                do_arrows(d.0, d.1, steps, false)
            } else if p.1 == 0 && n.0 == 3 {
                // must move right first
                do_arrows(d.0, d.1, steps, true)
            } else {
                // move in either direction
                std::cmp::min(
                    do_arrows(d.0, d.1, steps, true),
                    do_arrows(d.0, d.1, steps, false),
                )
            }
        })
        .sum::<usize>()
        * sequence[0..3].parse::<usize>().unwrap()
}

pub fn part_one(input: &str) -> Option<usize> {
    memoized_flush_do_arrows();
    Some(input.lines().map(|s| enter_sequence(s, 2)).sum())
}

pub fn part_two(input: &str) -> Option<usize> {
    memoized_flush_do_arrows();
    Some(input.lines().map(|s| enter_sequence(s, 25)).sum())
}