#! /usr/bin/env python3
import functools


TREE = "#"
SNOW = "."


class LoopingList(list):

    def __getitem__(self, key):
        if isinstance(key, int):
            return super().__getitem__(key % len(self))
        if isinstance(key, slice):
            return [
                self[x]
                for x in range(key.start, key.stop, key.step or 1)
            ]


def read_map():
    world_map = []
    with open("input.txt") as f:
        for line in f:
            world_map.append(LoopingList(line.strip()))
    return world_map


class World(object):

    def __init__(self, world_map, x=0, y=0):
        self._world_map = world_map
        self._x = x
        self._y = y
        # Default slopes
        self._dx = 1
        self._dy = 1

    def move(self, dx, dy, should_print=False, print_buffer=0):
        self._x += dx
        self._y += dy
        if should_print:
            print("".join(self._world_map[self._y][0:self._x+print_buffer]))
        return self._world_map[self._y][self._x]

    def print_row(self):
        print("".join(self._world_map[self._y]))

    def peek(self, dx, dy):
        return self._world_map[dy + self._y][dx + self._x]

    def at_end(self):
        return self._y >= len(self._world_map) - 1

    def set_slope(self, dx, dy):
        self._dx = dx
        self._dy = dy

    def ski(self, count_trees=True):
        trees = 0
        while not self.at_end():
            try:
                space = self.move(self._dx, self._dy)
                if count_trees and space == TREE:
                    trees += 1
            except IndexError:
                print("At end of the world")
        return trees

    def take_lift(self):
        self._x = 0
        self._y = 0


def part1():
    world_map = read_map()
    world = World(world_map)
    world.set_slope(3, 1)
    trees = world.ski()

    print("Total trees", trees)


def part2():
    slopes = {
        (1, 1): None,
        (3, 1): None,
        (5, 1): None,
        (7, 1): None,
        (1, 2): None,
    }

    world_map = read_map()
    world = World(world_map)

    for slope in slopes:
        world.take_lift()
        world.set_slope(*slope)
        slopes[slope] = world.ski()

    print(slopes)
    print("Answer:", functools.reduce(lambda x, y: x*y, slopes.values()))


if __name__ == "__main__":
    part1()
    part2()