#!/usr/bin/env python3

import sys

import numpy as np
import rich.progress

input_file = sys.argv[1]

with open(input_file) as fd:
    lines = [line.rstrip() for line in fd.readlines()]

buyers = len(lines)


def prng(secwet: int) -> int:
    secwet ^= secwet * 64
    secwet %= 16777216
    secwet ^= secwet // 32
    secwet %= 16777216
    secwet ^= secwet * 2048
    secwet %= 16777216
    return secwet


buyers_ban: list[list[int]] = []
buyers_dif: list[list[int]] = []

for line in rich.progress.track(lines, description="Calculating future banana prices"):
    secwet = int(line)
    buyer_ban: list[int] = []
    buyer_dif: list[int] = []
    last_ban = secwet % 10
    for i in range(2000):
        secwet = prng(secwet)
        ban = secwet % 10
        buyer_ban.append(ban)
        dif = ban - last_ban
        buyer_dif.append(dif)
        last_ban = ban
    #     print(f"{secwet=} {ban=} {dif=}")
    # print(f"{buyer_ban=}")
    # print(f"{buyer_dif=}")
    buyers_ban.append(buyer_ban)
    buyers_dif.append(buyer_dif)

buyers_dif_np = np.array(buyers_dif)
buyers_ban_np = np.array(buyers_ban)

sequence = tuple[int, int, int, int]


def totbans(seq: sequence) -> int:
    match = None
    for i, num in enumerate(seq):
        nmatch = buyers_dif_np == num
        if match is not None:
            # nmatch = sp.ndimage.shift(nmatch, (0, -i))
            nmatch = np.roll(nmatch, -i, axis=1)
            nmatch &= match
        match = nmatch

    # bans = buyers_ban_np * match
    # found = match.max(axis=1)
    # indexes = np.argmax(match > 0, axis=1)

    tot = 0
    assert match is not None
    for b, buyer_match in enumerate(match):
        if not buyer_match.max():
            continue
        arg: int = np.argmax(buyer_match > 0)
        try:
            ban = buyers_ban_np[b, arg+3]
            tot += ban
        except IndexError:
            pass # shrug
    return tot

    for buyer in range(buyers):
        buyer_dif = buyers_dif[buyer]
        for i in range(2000 - 4):
            if (
                buyer_dif[i] == seq[0]
                and buyer_dif[i + 1] == seq[1]
                and buyer_dif[i + 2] == seq[2]
                and buyer_dif[i + 3] == seq[3]
            ):
                # if tuple(buyer_dif[i : i + 4]) == seq:
                tot += buyers_ban[buyer][i + 3]
                break
    return tot


# print(f"{totbans((6, -1, -1, 0))=}") # demo0
# print(f"{totbans((-2, 1, -1, 3))=}") # demo2 aoc
# print(f"{totbans((6, -4, 4, -9))=}") # demo2 first

all_seqs: set[sequence] = set()

for buyer in rich.progress.track(
    range(buyers), description="Generating possible sequences"
):
    buyer_dif = buyers_dif[buyer]
    for i in range(2000 - 4):
        seq: sequence = tuple(buyer_dif[i : i + 4])
        all_seqs.add(seq)
print(f"{len(all_seqs)=}")

maxi = 0
max_seq = None
for seq in rich.progress.track(all_seqs, description="Finding score for sequences"):
    tb = totbans(seq)
    if tb > maxi:
        maxi = tb
        max_seq = seq


print(f"{max_seq=}")
print(maxi)
# 1909 too low