#!/usr/bin/env python3

from intcode import preproc
from itertools import permutations
from lib import last
from functools import cache

stack_size = 5
fst_amp_input = 0


def partI(amp):
    run_amp = cache(lambda *args : last(amp(*args)))

    phase_range = range(stack_size)
    best = 0 

    for perm in permutations(phase_range):
        amp_in = fst_amp_input
        for phase in perm:
            amp_in = run_amp(phase, amp_in)
        best = max(amp_in, best)
    return best
   
def partII(amp):
    phase_range = range(5, 10)
    best = 0 

    for perm in permutations(phase_range):
        head, *tail = map(amp, perm)
        amp_chain = sum(tail, head)

        amp_chain.append_input(fst_amp_input)
        
        for amp_in in amp_chain:
            amp_chain.write_input(amp_in)
        best = max(amp_in, best)
    return best

from main import Tests
tests = Tests()

tests.add(
    "3,15,3,16,1002,16,10,16,1,16,15,15,4,15,99,0,0", partI=43210)

tests.add(
    "3,23,3,24,1002,24,10,24,1002,23,-1,23," +
    "101,5,23,23,1,24,23,23,4,23,99,0,0", partI=54321)

tests.add(
    "3,31,3,32,1002,32,10,32,1001,31,-2,31,1007,31,0,33," +
    "1002,33,7,33,1,33,31,31,1,32,31,31,4,31,99,0,0,0", partI=65210)

tests.add(
    "3,26,1001,26,-4,26,3,27,1002,27,2,27,1,27,26," +
    "27,4,27,1001,28,-1,28,1005,28,6,99,0,0,5", partII=139629729)

tests.add(
    "3,52,1001,52,-5,52,3,53,1,52,56,54,1007,54,5,55,1005,55,26,1001,54," +
    "-5,54,1105,1,12,1,53,54,53,1008,54,0,55,1001,55,1,55,2,53,55,53,4," +
    "53,1001,56,-1,56,1005,56,6,99,0,0,0,0,10", partII=18216)