use aoc_runner_derive::{aoc, aoc_generator};
#[aoc_generator(day6, part1)]
fn parse(input: &str) -> Vec<Vec<i32>> {
    let input: Vec<Vec<i32>> = input.split('\n') // Separate the Time and Distance lines
        .map(|line| {
            line[line.find(':').unwrap() + 1..] // Drop "Time:" and "Distance:"
                .split_whitespace() // Split the numbers into their own elements.
                .map(|num| num.parse::<i32>().expect("Couldn't parse number")) // Parse numbers into i32
                .collect::<Vec<_>>()
        }).collect::<Vec<_>>(); // collect into Vec

        input
}

#[aoc(day6, part1)]
fn part1(input: &Vec<Vec<i32>>) -> i32 {
    let mut valid_total = 1;

    for round in 0..input.first().unwrap().len() {
        let time = input[0][round];
        let dist = input[1][round];
        let mut valid = 0;

        for remaining_time in 0..time {
            if (time - remaining_time) * remaining_time > dist {
                valid += 1;
            }
        }
        valid_total *= valid;
    }

    valid_total
}

#[aoc_generator(day6, part2)]
fn parse_part2(input: &str) -> Vec<i64> {
    let input: Vec<i64> = input.split('\n') // Separate the Time and Distance lines
        .map(|line| {
            line[line.find(':').unwrap() + 1..] // Drop "Time:" and "Distance:"
                .split_whitespace() // Split the numbers into their own elements
                .flat_map(|s| s.chars()).collect::<String>() // Combine the strings into a single one
                .parse::<i64>().expect("Couldn't parse number") // Parse numbers into i32
        }).collect(); // Collect into Vec

        input
}

#[aoc(day6, part2)]
fn part2(input: &Vec<i64>) -> i32 {
    let time = input[0];
    let dist = input[1];
    let mut valid = 0;

    for remaining_time in 0..time {
        if (time - remaining_time) * remaining_time > dist {
            valid += 1;
        }
    }

    valid
}


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

    const EX: &str = r"Time:      7  15   30
Distance:  9  40  200";

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

    #[test]
    fn part2_example() {
        assert_eq!(part2(&parse_part2(EX)), 71503);
    }
}