from tqdm import tqdm
from os.path import dirname,isfile,realpath,join

from collections.abc import Callable
def profiler(method) -> Callable[..., any]:
    from time import perf_counter_ns
    def wrapper_method(*args: any, **kwargs: any) -> any:
        start_time = perf_counter_ns()
        ret = method(*args, **kwargs)
        stop_time = perf_counter_ns() - start_time
        time_len = min(9, ((len(str(stop_time))-1)//3)*3)
        time_conversion = {9: 'seconds', 6: 'milliseconds', 3: 'microseconds', 0: 'nanoseconds'}
        print(f"Method {method.__name__} took : {stop_time / (10**time_len)} {time_conversion[time_len]}")
        return ret
    return wrapper_method

# Process a secret to a new secret number
# @param n: The secret number to be processed
# @return: The new secret number after processing
def process_secret(n: int) -> int:
    """ 
    Process a secret number by XORing it with the result of shifting left and right operations on itself. 
    The process involves several bitwise operations to ensure that the resulting number is different from the original one.
    
    First, multiply the original secret number by 64 with a left bit shift, then XOR the original secret number with the new number and prune to get a new secret number
    Second, divide the previous secret number by 32 with a right bit shift, then XOR the previous secret number with the new number and prune to get another new secret number
    Third, multiply the previous secret number by 2048 with a left bit shift, then XOR the previous secret number with the new number and prune to get the final secret number
    Finally, return the new secret number after these operations.

    """
    n ^= (n << 6) 
    n &= 0xFFFFFF
    n ^= (n >> 5)
    n &= 0xFFFFFF
    n ^= (n << 11)
    return n & 0xFFFFFF

# Solve Part 1 and Part 2 of the challenge at the same time
@profiler
def solve(secrets: list[int]) -> tuple[int, int]:
    # Build a dictionary for each buyer with the tuple of changes being the key and the sum of prices of the earliest occurrence for each buyer as the value
    # At the same time we solve Part 1 of the challenge by adding the last price for each secret
    last_price_sum = 0
    changes_map = {}
    for start_secret in (secrets):
        # Keep track of seen patterns for current secret
        changes_seen = set()
        # tuple of last 4 changes
        # first change is 0 because it is ignored
        last_four_changes = tuple([
                                    (cur_secret:=process_secret(start_secret)), 
                                   -(cur_secret%10) + ((cur_secret:=process_secret(cur_secret)) % 10) , 
                                   -(cur_secret%10) + ((cur_secret:=process_secret(cur_secret)) % 10) , 
                                   -(cur_secret%10) + ((cur_secret:=process_secret(cur_secret)) % 10) 
                                   ]
                                   )
        current_price = sum(last_four_changes)
        # Map 4-tuple of changes -> sum of prices index of earliest occurrence for all secrets
        for i in range(3, 1999):
            # sliding window of last four changes
            last_four_changes = (*last_four_changes[1:], -(cur_secret%10) + (current_price := (cur_secret:=process_secret(cur_secret)) % 10) )
            # if we have seen this pattern before, then we continue to next four changes
            # otherwise, we add the price to the mapped value
            # this ensures we only add the first occurence of a patten for each list of prices each secret produces
            if last_four_changes not in changes_seen:
                # add to the set of seen patterns for this buyer
                changes_seen.add(last_four_changes)
                # If not recorded yet, store the price
                # i+4 is the index of the price where we sell
                changes_map[last_four_changes] = changes_map.get(last_four_changes, 0) + current_price

        # Sum the 2000th price to the total sum for all secrets
        last_price_sum += cur_secret
    # Return the sum of all prices at the 2000th iteration and the maximum amount of bananas that one pattern can obtain
    return last_price_sum,max(changes_map.values())

if __name__ == "__main__":
    # Read buyers' initial secrets from file or define them here
    BASE_DIR = dirname(realpath(__file__))
    with open(join(BASE_DIR, r'input'), "r") as f:
        secrets = [int(x) for x in f.read().split()]
    last_price_sum,best_score = solve(secrets)
    print("Part 1:")
    print(f"sum of each of the 2000th secret number:", last_price_sum)
    print("Part 2:")
    print("Max bananas for one of patterns:", best_score)