from aoc import AOCDay
from typing import Any, List


def follow_directions(path: List) -> (int, int, int):
    pos = 0
    depth = 0
    aim = 0
    for direction in path:
        if direction[0] == 'forward':
            pos += direction[1]
            depth += aim * direction[1]
        elif direction[0] == 'down':
            aim += direction[1]
        elif direction[0] == 'up':
            aim -= direction[1]

    return pos, depth, aim


class Day(AOCDay):
    test_solutions_p1 = [150]
    test_solutions_p2 = [900]

    def part1(self) -> Any:
        path = self.getInputAsArraySplit(" ", [str, int])
        pos, _, depth = follow_directions(path)

        return pos * depth

    def part2(self) -> Any:
        path = self.getInputAsArraySplit(" ", [str, int])
        pos, depth, _ = follow_directions(path)

        return pos * depth