import itertools
from intcode import IntCode
from tools.aoc import AOCDay
from typing import Any, List


def get_comps(memory: List[int]) -> List[IntCode]:
    comps = []
    for _ in range(5):
        comps.append(IntCode(memory.copy()))

    return comps


class Day(AOCDay):
    test_solutions_p1 = [43210, 54321, 65210]
    test_solutions_p2 = [139629729, 18216]

    def part1(self) -> Any:
        init_memory = self.getInputAsArraySplit(",", int)
        max_signal = 0
        phase_settings = [0, 1, 2, 3, 4]
        for s in itertools.permutations(phase_settings):
            comps = get_comps(init_memory)
            for i, input1 in enumerate(s):
                if i == 0:
                    input2 = 0
                else:
                    input2 = comps[i-1].getOutput()

                comps[i].addInput(input1)
                comps[i].addInput(input2)
                comps[i].run()

            signal = comps[4].getOutput()
            if signal > max_signal:
                max_signal = signal

        return max_signal

    def part2(self) -> Any:
        init_memory = self.getInputAsArraySplit(",", int)
        max_signal = 0
        phase_settings = [5, 6, 7, 8, 9]
        for s in itertools.permutations(phase_settings):
            comps = get_comps(init_memory.copy())

            for i, input1 in enumerate(s):
                if i == 0 and not comps[i].isRunning():
                    input2 = 0
                else:
                    input2 = comps[i-1].getOutput()

                comps[i].addInput(input1)
                comps[i].addInput(input2)
                comps[i].start()

            while not comps[4].isHalted():
                for i in range(5):
                    if not comps[i].isHalted() and not comps[i-1].isHalted():
                        comps[i].addInput(comps[i-1].getOutput())

            signal = comps[4].getOutput()
            if signal > max_signal:
                max_signal = signal

            for i in range(5):
                comps[i].join()

        return max_signal