from tools.aoc import AOCDay
from typing import Any


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):
    test_solutions_p1 = [5934]
    test_solutions_p2 = [26984457539]

    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))