from collections import defaultdict
from tools.aoc import AOCDay
from typing import Any


def get_next_secret(secret: int, steps: int = 1) -> int:
    for _ in range(steps):
        new_secret = secret * 64
        secret = (secret ^ new_secret) % 16777216
        new_secret = secret // 32
        secret = (secret ^ new_secret) % 16777216
        new_secret = secret * 2048
        secret = (secret ^ new_secret) % 16777216

    return secret


class Day(AOCDay):
    inputs = [
        [
            (37327623, "input22_test"),
            (14869099597, "input22"),
        ],
        [
            (23, "input22_test2"),
            (1717, "input22"),
        ],
    ]

    def parse_input(self) -> list[int]:
        return self.getInput(return_type=int)

    def part1(self) -> Any:
        return sum(get_next_secret(x, 2000) for x in self.parse_input())

    def part2(self) -> Any:
        sequences = defaultdict(int)
        max_bananas = 0
        for secret in self.parse_input():
            sequence = (-10, -10, -10, -10)
            seen = set()
            for i in range(2000):
                next_secret = get_next_secret(secret)
                price = next_secret % 10
                diff = price - secret % 10
                secret = next_secret
                sequence = sequence[-3:] + (diff,)
                if sequence in seen or i < 4:
                    continue
                seen.add(sequence)
                sequences[sequence] += price
                if sequences[sequence] > max_bananas:
                    max_bananas = sequences[sequence]

        return max_bananas


if __name__ == "__main__":
    day = Day(2024, 22)
    day.run(verbose=True)