fn aoc_06_part01(input: &[Vec<usize>], days: usize) -> usize {
    // We only care about the last day technically...
    let mut last_day = input[input.len() - 1].clone();
    let mut to_add: usize;
    for _day in 0..(days - input.len() + 1) {
        to_add = 0;
        for lanternfish in last_day.iter_mut() {
            if lanternfish == &0 {
                *lanternfish = 6;
                to_add += 1;
            } else {
                *lanternfish -= 1;
            }
        }
        for _add in 0..to_add {
            last_day.push(8);
        }
    }

    last_day.len()
}

fn aoc_06_part02(input: &[Vec<usize>], days: usize) -> usize {
    // Basically we dont care about individual values, all we care is each number of fish per day,
    // which is a fixed vec...
    let last_day = input[input.len() - 1].clone();
    let mut aggregated_lanternfish: Vec<usize> = last_day.iter().fold(vec![0; 9], |mut acc, val| {
        acc[*val] += 1;
        acc
    });

    let mut to_add: usize;
    for _day in 0..(days - input.len() + 1) {
        // Very simple, each 0 goes back to 6
        // Each 0 creates one 8 each...
        to_add = aggregated_lanternfish[0];

        // Depleting each
        for i in 1..9 {
            aggregated_lanternfish[i - 1] = aggregated_lanternfish[i];
        }

        aggregated_lanternfish[6] += to_add;
        aggregated_lanternfish[8] = to_add;
    }

    aggregated_lanternfish.iter().sum()
}

fn main() {
    let s = include_str!("../input.txt").lines()
        // .map(|s| s.split_once(": ").unwrap_or(("", "")).1)
        .map(|s| s.split(',').filter_map(|s| s.parse::<usize>().ok()).collect::<Vec<usize>>())
        .collect::<Vec<Vec<usize>>>();

    let result = aoc_06_part01(&s, 80);
    println!("Number of lanternfish after 80 days: {:?}", result);


    let result = aoc_06_part02(&s, 256);
    println!("Number of lanternfish after 256 days: {:?}", result);
}