use aoc_runner_derive::{aoc, aoc_generator};

#[aoc_generator(day3)]
fn parse(input: &str) -> Vec<Vec<char>> {
    let input = input.lines().map(|line| {
        line.chars().collect::<Vec<_>>()
    }).collect::<Vec<_>>();

    input
}

#[aoc(day3, part1)]
fn part1(input: &Vec<Vec<char>>) -> u64 {
    count_trees(1, 3, input)
}

#[aoc(day3, part2)]
fn part2(input: &Vec<Vec<char>>) -> u64 {
    vec![(1,1), (1,3), (1,5), (1,7), (2,1)].iter()
        .map(|(rise, run)| {
            count_trees(*rise, *run, input)
        }).product()
}

fn count_trees(rise: usize, run: usize, input: &Vec<Vec<char>>) -> u64 {
    let mut x = 0;
    let mut tree_count = 0;

    for y in (rise..input.len()).step_by(rise) {
        x += run;
        if input[y][x % input[0].len()] == '#' {
            tree_count += 1;
        }
    }
    tree_count
}


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

const EX: &str = r"..##.......
#...#...#..
.#....#..#.
..#.#...#.#
.#...##..#.
..#.##.....
.#.#.#....#
.#........#
#.##...#...
#...##....#
.#..#...#.#";

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

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