use aoc_runner_derive::{aoc, aoc_generator};

#[derive(Debug, PartialEq, Eq)]
enum NumButton {
    One,
    Two,
    Three,
    Four,
    Five,
    Six,
    Seven,
    Eight,
    Nine,
    Zero,
    A,
}

impl NumButton {
    const fn pos(&self) -> (usize, usize) {
        match self {
            NumButton::One => (0, 2),
            NumButton::Two => (1, 2),
            NumButton::Three => (2, 2),
            NumButton::Four => (0, 1),
            NumButton::Five => (1, 1),
            NumButton::Six => (2, 1),
            NumButton::Seven => (0, 0),
            NumButton::Eight => (1, 0),
            NumButton::Nine => (2, 0),
            NumButton::Zero => (1, 3),
            NumButton::A => (2, 3),
        }
    }

    const fn index(&self) -> usize {
        self.pos().1 * 3 + self.pos().0
    }

    const fn value(&self) -> usize {
        match self {
            NumButton::One => 1,
            NumButton::Two => 2,
            NumButton::Three => 3,
            NumButton::Four => 4,
            NumButton::Five => 5,
            NumButton::Six => 6,
            NumButton::Seven => 7,
            NumButton::Eight => 8,
            NumButton::Nine => 9,
            NumButton::Zero => 0,
            NumButton::A => unimplemented!(),
        }
    }

    const fn from_index(idx: usize) -> Self {
        match idx {
            0 => NumButton::Seven,
            1 => NumButton::Eight,
            2 => NumButton::Nine,
            3 => NumButton::Four,
            4 => NumButton::Five,
            5 => NumButton::Six,
            6 => NumButton::One,
            7 => NumButton::Two,
            8 => NumButton::Three,
            10 => NumButton::Zero,
            11 => NumButton::A,
            _ => unreachable!(),
        }
    }

    const fn paths_to(&self, other: &NumButton) -> [Path; 2] {
        let start = self.pos();
        let end = other.pos();

        let bx = if start.0 < end.0 { RIGHT } else { LEFT };
        let dx = start.0.abs_diff(end.0);

        let by = if start.1 < end.1 { DOWN } else { UP };
        let dy = start.1.abs_diff(end.1);

        if (start.0 == 0 && end.1 == 3) || start.1 == end.1 {
            let mut x = [DistanceSize::MAX; 6];
            let y = [DistanceSize::MAX; 6];
            let mut i = dx;
            while i > 0 {
                x[i - 1] = bx;
                i -= 1;
            }
            let mut i = dy;
            while i > 0 {
                x[i - 1 + dx] = by;
                i -= 1;
            }
            x[dx + dy] = A;
            [x, y]
        } else if (end.0 == 0 && start.1 == 3) || start.0 == end.0 {
            let x = [DistanceSize::MAX; 6];
            let mut y = [DistanceSize::MAX; 6];
            let mut i = dx;
            while i > 0 {
                y[i - 1 + dy] = bx;
                i -= 1;
            }
            let mut i = dy;
            while i > 0 {
                y[i - 1] = by;
                i -= 1;
            }
            y[dx + dy] = A;
            [y, x]
        } else {
            let mut x = [DistanceSize::MAX; 6];
            let mut y = [DistanceSize::MAX; 6];
            let mut i = dx;
            while i > 0 {
                x[i - 1] = bx;
                y[i - 1 + dy] = bx;
                i -= 1;
            }
            let mut i = dy;
            while i > 0 {
                x[i - 1 + dx] = by;
                y[i - 1] = by;
                i -= 1;
            }
            x[dx + dy] = A;
            y[dx + dy] = A;
            [x, y]
        }
    }
}

type Input = String;

#[aoc_generator(day21)]
fn parse(input: &str) -> Input {
    input.to_string()
}

type Path = [DistanceSize; 6];

type DistanceSize = usize;
type DistanceMatrix = [[DistanceSize; 5]; 5];

const fn shortest_dir_path(path: &Path, steps: &DistanceMatrix) -> DistanceSize {
    let mut sum = 0;
    let mut pos = A;
    let mut i = 0;

    while i < 6 && path[i] != DistanceSize::MAX {
        sum += steps[pos][path[i]];
        pos = path[i];
        i += 1;
    }

    sum
}

const fn shortest_path(
    path: &[usize],
    paths: &[[[Path; 2]; 12]; 12],
    steps: &DistanceMatrix,
) -> DistanceSize {
    let mut sum = 0;
    let mut pos = NumButton::A.index();
    let mut i = 0;
    while i < path.len() {
        let step_paths = paths[pos][path[i]];
        sum += if step_paths[1][0] != DistanceSize::MAX {
            min(
                shortest_dir_path(&step_paths[0], steps),
                shortest_dir_path(&step_paths[1], steps),
            )
        } else {
            shortest_dir_path(&step_paths[0], steps)
        };
        pos = path[i];
        i += 1;
    }
    sum
}

const UP: usize = 0;
const DOWN: usize = 1;
const LEFT: usize = 2;
const RIGHT: usize = 3;
const A: usize = 4;

const fn min(a: usize, b: usize) -> usize {
    if a < b {
        a
    } else {
        b
    }
}

const fn precompute_steps_new<const N: usize>(matrix: DistanceMatrix) -> DistanceMatrix {
    let mut steps = 0;
    let mut matrix = matrix;

    while steps < N {
        let up = matrix[A][UP] + matrix[UP][A];
        let down = matrix[A][DOWN] + matrix[DOWN][A];
        let left = matrix[A][LEFT] + matrix[LEFT][A];
        let right = matrix[A][RIGHT] + matrix[RIGHT][A];
        let left_left = matrix[A][LEFT] + matrix[LEFT][LEFT] + matrix[LEFT][A];
        let right_right = matrix[A][RIGHT] + matrix[RIGHT][RIGHT] + matrix[RIGHT][A];
        let left_down = matrix[A][LEFT] + matrix[LEFT][DOWN] + matrix[DOWN][A];
        let down_left = matrix[A][DOWN] + matrix[DOWN][LEFT] + matrix[LEFT][A];
        let left_up = matrix[A][LEFT] + matrix[LEFT][UP] + matrix[UP][A];
        let up_left = matrix[A][UP] + matrix[UP][LEFT] + matrix[LEFT][A];
        let right_down = matrix[A][RIGHT] + matrix[RIGHT][DOWN] + matrix[DOWN][A];
        let down_right = matrix[A][DOWN] + matrix[DOWN][RIGHT] + matrix[RIGHT][A];
        let right_up = matrix[A][RIGHT] + matrix[RIGHT][UP] + matrix[UP][A];
        let up_right = matrix[A][UP] + matrix[UP][RIGHT] + matrix[RIGHT][A];
        let down_left_left =
            matrix[A][DOWN] + matrix[DOWN][LEFT] + matrix[LEFT][LEFT] + matrix[LEFT][A];
        let right_right_up =
            matrix[A][RIGHT] + matrix[RIGHT][RIGHT] + matrix[RIGHT][UP] + matrix[UP][A];

        matrix[UP][UP] = 1;
        matrix[UP][DOWN] = down;
        matrix[UP][LEFT] = down_left;
        matrix[UP][RIGHT] = min(down_right, right_down);
        matrix[UP][A] = right;

        matrix[DOWN][UP] = up;
        matrix[DOWN][DOWN] = 1;
        matrix[DOWN][LEFT] = left;
        matrix[DOWN][RIGHT] = right;
        matrix[DOWN][A] = min(up_right, right_up);

        matrix[LEFT][UP] = right_up;
        matrix[LEFT][DOWN] = right;
        matrix[LEFT][LEFT] = 1;
        matrix[LEFT][RIGHT] = right_right;
        matrix[LEFT][A] = right_right_up;

        matrix[RIGHT][UP] = min(up_left, left_up);
        matrix[RIGHT][DOWN] = left;
        matrix[RIGHT][LEFT] = left_left;
        matrix[RIGHT][RIGHT] = 1;
        matrix[RIGHT][A] = up;

        matrix[A][UP] = left;
        matrix[A][DOWN] = min(left_down, down_left);
        matrix[A][LEFT] = down_left_left;
        matrix[A][RIGHT] = down;
        matrix[A][A] = 1;

        steps += 1;
    }

    matrix
}

const fn precompute_paths() -> [[[Path; 2]; 12]; 12] {
    let mut paths = [[[[DistanceSize::MAX; 6], [DistanceSize::MAX; 6]]; 12]; 12];

    let mut start = 0usize;
    while start < 12 {
        let mut end = 0usize;
        if start == 9 {
            start += 1;
            continue;
        }
        while end < 12 {
            if end == 9 {
                end += 1;
                continue;
            }

            paths[start][end] = NumButton::from_index(start).paths_to(&NumButton::from_index(end));

            end += 1;
        }
        start += 1;
    }

    paths
}

const fn precompute_num_paths(
    steps: &[[usize; 5]; 5],
    paths: &[[[[usize; 6]; 2]; 12]; 12],
) -> [[[usize; 12]; 12]; 12] {
    let mut matrix = [[[0; 12]; 12]; 12];

    let mut one = 0usize;
    while one < 12 {
        if one == 9 || one == 11 {
            one += 1;
            continue;
        }
        let mut two = 0usize;
        while two < 12 {
            if two == 9 || two == 11 {
                two += 1;
                continue;
            }
            let mut three = 0usize;
            while three < 12 {
                if three == 9 || three == 11 {
                    three += 1;
                    continue;
                }

                let num = NumButton::from_index(one).value() * 100
                    + NumButton::from_index(two).value() * 10
                    + NumButton::from_index(three).value();
                matrix[one][two][three] =
                    num * shortest_path(&[one, two, three, NumButton::A.index()], paths, steps);

                three += 1;
            }
            two += 1;
        }
        one += 1;
    }

    matrix
}

const CHAR_IDX: [usize; 10] = [
    NumButton::Zero.index(),
    NumButton::One.index(),
    NumButton::Two.index(),
    NumButton::Three.index(),
    NumButton::Four.index(),
    NumButton::Five.index(),
    NumButton::Six.index(),
    NumButton::Seven.index(),
    NumButton::Eight.index(),
    NumButton::Nine.index(),
];

#[aoc(day21, part1)]
fn part1(input: &Input) -> DistanceSize {
    let num_paths = const {
        let steps = precompute_steps_new::<2>([[1; 5]; 5]);
        let paths = precompute_paths();
        precompute_num_paths(&steps, &paths)
    };

    input
        .as_bytes()
        .chunks(5)
        .map(|chunk| {
            let num = [
                CHAR_IDX[chunk[0] as usize - 48],
                CHAR_IDX[chunk[1] as usize - 48],
                CHAR_IDX[chunk[2] as usize - 48],
            ];
            num_paths[num[0]][num[1]][num[2]]
        })
        .sum()
}

#[aoc(day21, part2)]
fn part2(input: &Input) -> DistanceSize {
    let num_paths = const {
        let steps = precompute_steps_new::<25>([[1; 5]; 5]);
        let paths = precompute_paths();
        precompute_num_paths(&steps, &paths)
    };

    input
        .as_bytes()
        .chunks(5)
        .map(|chunk| {
            let num = [
                CHAR_IDX[chunk[0] as usize - 48],
                CHAR_IDX[chunk[1] as usize - 48],
                CHAR_IDX[chunk[2] as usize - 48],
            ];
            num_paths[num[0]][num[1]][num[2]]
        })
        .sum()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn part1_example() {
        assert_eq!(part1(&parse("379A")), 64 * 379);
        assert_eq!(part1(&parse("179A")), 68 * 179);
        assert_eq!(part1(&parse("456A")), 64 * 456);
        assert_eq!(part1(&parse("980A")), 60 * 980);
        assert_eq!(part1(&parse("029A")), 68 * 29);
    }
}