from tools.aoc import AOCDay
from typing import Any


class Day(AOCDay):
    inputs = [
        [
            (5, "input6_test"),
            (12841, "input6")
        ],
        [
            (4, "input6_test"),
            (8038, "input6")
        ]
    ]

    def distribute(self, p2: bool = False) -> int:
        bank = self.getInputAsArraySplit("\t", int)
        count = 0
        seen = [",".join(map(str, bank))]
        while True:
            count += 1
            i, c = bank.index(max(bank)), max(bank)
            bank[i] = 0
            while c > 0:
                i += 1
                bank[i%len(bank)] += 1
                c -= 1
            seq = ",".join(map(str, bank))
            if seq in seen:
                if p2:
                    return count - seen.index(seq)
                else:
                    return count
            else:
                seen.append(seq)

    def part1(self) -> Any:
        return self.distribute()

    def part2(self) -> Any:
        return self.distribute(True)


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