#!/usr/bin/env python3

import collections
import math
import sys

import numpy as np
import rich.progress

dtype = np.int32

input_file = sys.argv[1]

with open(input_file) as fd:
    buyers_init_secret = np.array([int(line) for line in fd.readlines()], dtype=dtype)

buyers = len(buyers_init_secret)

ITERATIONS = 2000
SEQUENCE_LENGTH = 4


buyers_prices = np.ndarray((buyers, ITERATIONS), dtype=dtype)
buyers_diffs = np.ndarray((buyers, ITERATIONS), dtype=dtype)

buyers_secret = buyers_init_secret.copy()
last_buyers_price = buyers_init_secret % 10
for i in rich.progress.track(range(ITERATIONS), description="Simulating stock market"):
    buyers_secret ^= buyers_secret * 64
    buyers_secret %= 16777216
    buyers_secret ^= buyers_secret // 32
    buyers_secret %= 16777216
    buyers_secret ^= buyers_secret * 2048
    buyers_secret %= 16777216
    buyers_price = buyers_secret % 10
    buyers_diff = buyers_price - last_buyers_price
    buyers_prices[:, i] = buyers_price
    buyers_diffs[:, i] = buyers_diff
    last_buyers_price = buyers_price

# Compress sequence tuples into a single integer

# Move to positive numbers
trans = -buyers_diffs.min()
buyers_diffs_translated = buyers_diffs + trans
# Decide on a value to shift
shift = math.ceil(math.log2(buyers_diffs_translated.max()))

buyers_sequences = buyers_diffs_translated.copy()
for i in range(1, SEQUENCE_LENGTH):
    buyers_sequences += np.roll(buyers_diffs_translated << (shift * i), i, axis=1)
# Make first few sequences invalid
buyers_sequences[:, :SEQUENCE_LENGTH] = -1

# NEW (knowing it's best to iterate per buyer than per sequence) ~ 5 seconds
# Inspired by the following, which is even faster, probably because data locality
# achieves more than SIMD-fying?
# https://github.com/mkern75/AdventOfCodePython/blob/23b6becdc873c6b865e783122a7dbce0b5f40f60/year2024/Day22.py

max_sequence = 1 << shift * SEQUENCE_LENGTH
sequences_total = np.zeros(max_sequence, dtype=dtype)
seen = np.zeros((buyers, max_sequence), dtype=bool)

for b in rich.progress.track(range(buyers), description="Evaluating each buyer"):
    for i in range(SEQUENCE_LENGTH, ITERATIONS):
        seq = buyers_sequences[b, i]

        if seen[b, seq]:
            continue
        seen[b, seq] = True

        sequences_total[seq] += buyers_prices[b, i]

print(f"{sequences_total.argmax()=}")
print(sequences_total.max())

sys.exit(0)

# OLD (knowing you can compresses sequences only) ~ 1.5 minute


def totbans(seq: int) -> int:
    match = buyers_sequences == seq
    found = match.max(axis=1)
    indexes = np.argmax(match, axis=1)
    bans = buyers_prices[range(buyers), indexes]
    bans *= found
    return bans.sum()


def seq_to_int(seq: tuple[int, ...]) -> int:
    tot = 0
    for s, num in enumerate(seq):
        tot += (num + trans) << (SEQUENCE_LENGTH - s - 1) * shift
    return tot


print(f"{totbans(seq_to_int((-2, 1, -1, 3)))=}")

all_seqs: set[int] = set(buyers_sequences.flat) - {-1}

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)