from intcode import IntCode
import itertools
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):
    inputs = [
        [
            (43210, "test_input07_1_0"),
            (54321, "test_input07_1_1"),
            (65210, "test_input07_1_2"),
            (199988, "input07")
        ],
        [
            (139629729, "test_input07_2_0"),
            (18216, "test_input07_2_1"),
            (17519904, "input07")
        ]
    ]

    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


if __name__ == '__main__':
    day = Day(2019, 7)
    day.run(verbose=True)