#!/usr/bin/env python3
import aoclib

DAY = 3


def check_for_trees(right_step, down_step, treelines):
    x = 0
    y = 0
    tree_count = 0
    for treeline in treelines:
        if y % down_step != 0:
            y = y + 1
            continue

        while len(treeline) < x + 1:
            treeline = treeline + treeline

        if treeline[x] == '#':
            tree_count = tree_count + 1

        x = x + right_step
        y = y + 1

    return tree_count


def part1(test_mode=False):
    my_input = aoclib.getInputAsArray(day=3, test=test_mode)
    return check_for_trees(3, 1, my_input)


def part2(test_mode=False):
    my_input = aoclib.getInputAsArray(day=3, test=test_mode)
    tree_count1 = check_for_trees(1, 1, my_input)
    tree_count2 = check_for_trees(3, 1, my_input)
    tree_count3 = check_for_trees(5, 1, my_input)
    tree_count4 = check_for_trees(7, 1, my_input)
    tree_count5 = check_for_trees(1, 2, my_input)
    return tree_count1 * tree_count2 * tree_count3 * tree_count4 * tree_count5


if __name__ == '__main__':
    assert part1(test_mode=True) == 7, "Part 1 TEST FAILED"
    aoclib.printSolution(DAY, 1, part1())
    assert part2(test_mode=True) == 336, "Part 2 TEST FAILED"
    aoclib.printSolution(DAY, 2, part2())