use aoc_runner_derive::{aoc, aoc_generator};

type Input = Vec<usize>;

#[aoc_generator(day22)]
fn parse(input: &str) -> Input {
    input.lines().map(|n| n.parse().unwrap()).collect()
}

#[inline(always)]
fn next_number(number: usize) -> usize {
    let step1 = (number ^ (number << 6)) % (1 << 24);
    let step2 = (step1 ^ (step1 >> 5)) % (1 << 24);
    (step2 ^ step2 << 11) % (1 << 24)
}

#[inline(always)]
fn step_2000(number: usize) -> usize {
    let mut number = number;
    for _ in 0..2000 {
        number = next_number(number);
    }
    number
}

#[aoc(day22, part1)]
fn part1(input: &Input) -> usize {
    input.iter().map(|n| step_2000(*n)).sum()
}

#[aoc(day22, part2)]
fn part2(input: &Input) -> usize {
    let mut total = Box::new([[[[0; 19]; 19]; 19]; 19]);
    let mut max_total = 0;
    let mut last_monkey = Box::new([[[[0; 19]; 19]; 19]; 19]);
    for (monkey, secret) in input.iter().enumerate() {
        let mut one = *secret;
        let mut two = next_number(one);
        let mut three = next_number(two);
        let mut four = next_number(three);
        for _ in 0..2000 - 3 {
            let next = next_number(four);

            let idx_1 = (two % 10) - (one % 10) + 9;
            let idx_2 = (three % 10) - (two % 10) + 9;
            let idx_3 = (four % 10) - (three % 10) + 9;
            let idx_4 = (next % 10) - (four % 10) + 9;

            if last_monkey[idx_1][idx_2][idx_3][idx_4] <= monkey {
                last_monkey[idx_1][idx_2][idx_3][idx_4] = monkey + 1;
                total[idx_1][idx_2][idx_3][idx_4] += next % 10;
                max_total = max_total.max(total[idx_1][idx_2][idx_3][idx_4]);
            }

            one = two;
            two = three;
            three = four;
            four = next;
        }
    }
    max_total
}

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

    #[test]
    fn part1_example() {
        assert_eq!(part1(&parse("1")), 8685429);
        assert_eq!(part1(&parse("10")), 4700978);
        assert_eq!(part1(&parse("100")), 15273692);
        assert_eq!(part1(&parse("2024")), 8667524);
    }
}