use aoc_runner_derive::{aoc, aoc_generator};
use memoize::memoize;

#[aoc_generator(day3)]
fn parse(input: &str) -> Vec<Vec<u8>> {
    input
        .lines()
        .map(|bank| bank.bytes().map(|c| c - b'0').collect())
        .collect()
}

#[memoize]
fn max_joltage(bank: Vec<u8>, n: usize) -> u64 {
    if n == 1 {
        return *bank.iter().max().unwrap() as u64;
    }

    (0..bank.len() - n + 1)
        .map(|start| {
            bank[start] as u64 * 10u64.pow(n as u32 - 1)
                + max_joltage(bank[start + 1..].to_vec(), n - 1)
        })
        .max()
        .unwrap()

    // let mut max = 0;
    // for start in 0..bank.len() - n + 1 {
    //     let cur = bank[start] as u64 * 10u64.pow(n as u32 - 1)
    //         + *bank[start + 1..].iter().max().unwrap() as u64;
    //     if cur > max {
    //         max = cur
    //     }
    // }
    // max
}

#[aoc(day3, part1)]
fn part1(input: &[Vec<u8>]) -> u64 {
    input.iter().map(|bank| max_joltage(bank.clone(), 2)).sum()
}

#[aoc(day3, part2)]
fn part2(input: &[Vec<u8>]) -> u64 {
    input.iter().map(|bank| max_joltage(bank.clone(), 12)).sum()
}

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

    const EXAMPLE: &str = "987654321111111
811111111111119
234234234234278
818181911112111";

    #[test]
    fn part1_example() {
        assert_eq!(part1(&parse(EXAMPLE)), 357);
    }

    #[test]
    fn part2_example() {
        assert_eq!(part2(&parse(EXAMPLE)), 3121910778619);
    }
}