from tools.aoc import AOCDay
from typing import Any, List


def evolve_fishes(start_state: List[int], days: int) -> List[int]:
    fishes = [0 for _ in range(9)]
    for x in start_state:
        fishes[x] += 1

    for day in range(days):
        new_fishes = fishes[0]
        for x in range(1, 9):
            fishes[x - 1] = fishes[x]

        fishes[6] += new_fishes
        fishes[8] = new_fishes

    return fishes


class Day(AOCDay):
    inputs = [
        [
            (5934, "test_input06"),
            (346063, "input06")
        ],
        [
            (26984457539, "test_input06"),
            (1572358335990, "input06")
        ]
    ]

    def part1(self) -> Any:
        start_state = self.getInputAsArraySplit(",", int)
        return sum(evolve_fishes(start_state, 80))

    def part2(self) -> Any:
        start_state = self.getInputAsArraySplit(",", int)
        return sum(evolve_fishes(start_state, 256))


if __name__ == '__main__':
    day = Day(6)
    day.run(verbose=True)