AOC 2021 Day 6

# aoc202106.py

import pathlib
import sys


def parse(puzzle_input: str) -> list[int]:
    """ Parse input """
    return [int(char) for char in puzzle_input.split(',')]


def grow_shoal(initial_state: list[int], cycle_length: int,
               num_days: int) -> list[int]:
    """ Increase shoal of fish by num of days """
    stages = [0] * cycle_length
    for fish in initial_state:
        stages[fish] += 1
    for _ in range(num_days):
        zero_day = stages[0]
        stages = stages[1:] + stages[:1]
        stages[6] += zero_day
    return stages


def part1(data: list[int]) -> int:
    """ Solve part 1 """
    days = 80
    generation = 9
    return sum(grow_shoal(data, generation, days))


def part2(data: list[int]) -> int:
    """ Solve part 2 """
    days = 256
    generation = 9
    return sum(grow_shoal(data, generation, days))


def solve(puzzle_input: str) -> tuple[int, int]:
    """ Solve the puzzle for the given input """
    data = parse(puzzle_input)
    solution1 = part1(data)  # Correct answer was 391888 (with my data)
    solution2 = part2(data)  # Correct answer was 1754597645339 (with my data)

    return solution1, solution2


if __name__ == "__main__":
    for path in sys.argv[1:]:
        print(f"{path}:")
        puzzle_input = pathlib.Path(path).read_text().strip()
        solutions = solve(puzzle_input)
        print('\n'.join(str(solution) for solution in solutions))