defmodule Day3 do
  @example """
  ..##.......
  #...#...#..
  .#....#..#.
  ..#.#...#.#
  .#...##..#.
  ..#.##.....
  .#.#.#....#
  .#........#
  #.##...#...
  #...##....#
  .#..#...#.#
  """

  def part1(example \\ false) do
    count_trees(example, {3, 1})
  end

  def part2(example \\ false) do
    one_one = count_trees(example, {1, 1})
    three_one = count_trees(example, {3, 1})
    five_one = count_trees(example, {5, 1})
    seven_one = count_trees(example, {7, 1})
    one_two = count_trees(example, {1, 2})

    one_one * three_one * five_one * seven_one * one_two
  end

  def count_trees(example, slope) do
    example
    |> parse_input()
    |> count_trees_for_slope(slope)
  end

  def parse_input(example) do
    if(example, do: @example, else: File.read!("lib/day3/input.txt"))
    |> String.trim()
    |> String.split("\n")
    |> Enum.map(fn line ->
      line
      |> String.codepoints()
      |> Enum.map(fn
        "#" -> 1
        _ -> 0
      end)
      |> Stream.cycle()
    end)
  end

  def count_trees_for_slope(lines, slope) do
    do_count(lines, slope, 0, 0)
  end

  def do_count([current | rest], {dx, dy} = slope, x, trees) do
    rest = Enum.drop(rest, dy - 1)
    do_count(rest, slope, x + dx, trees + Enum.at(current, x))
  end

  def do_count([], _, _, trees) do
    trees
  end
end