import kotlinx.collections.immutable.PersistentList
import kotlinx.collections.immutable.persistentListOf
import kotlinx.collections.immutable.toPersistentList
import org.assertj.core.api.Assertions.assertThat
import org.junit.Test

class Day7AmplificationCircuit {


    private fun evaluateThrustersPower(inputs: List<Int>, program: PersistentList<Int>): Int {
        return inputs.fold(0) { power, input ->
            intCodeComputer(program, input = *intArrayOf(input, power))
                .runToHalt()
                .out
                .first()
                .toInt()
        }
    }

    private fun evaluateThrustersPowerWithFeedbackLoop(inputs: List<Int>, program: PersistentList<Int>): Int {

        val seed = listOf(
            intCodeComputer(program, input = *intArrayOf(inputs[0])),
            intCodeComputer(program, input = *intArrayOf(inputs[1])),
            intCodeComputer(program, input = *intArrayOf(inputs[2])),
            intCodeComputer(program, input = *intArrayOf(inputs[3])),
            intCodeComputer(program, input = *intArrayOf(inputs[4]))
        )

        return generateSequence(0 to seed) { (feed, amplifiers) ->
            try {
                val a = amplifiers[0].addInput { it + feed }.runToOutputOrHalt()
                val b = amplifiers[1].addInput { it + a.out.first() }.runToOutputOrHalt()
                val c = amplifiers[2].addInput { it + b.out.first() }.runToOutputOrHalt()
                val d = amplifiers[3].addInput { it + c.out.first() }.runToOutputOrHalt()
                val e = amplifiers[4].addInput { it + d.out.first() }.runToOutputOrHalt()
                e.out.first().toInt() to listOf(a, b, c, d, e)
            } catch (e: Exception) {
                if (e.message == "halted") {
                    null
                } else {
                    throw e
                }
            }
        }.takeWhile { (feed, amplifiers) -> amplifiers.all { it.pc >= 0 } }
            .last()
            .second
            .last()
            .out
            .first()
            .toInt()
    }

    @Test
    fun verifyTestData() {
        examples.forEach { (max, inputs, program) ->
            assertThat(evaluateThrustersPower(inputs, program)).isEqualTo(max)
        }
    }

    @Test
    fun verifyTestDataGold() {
        examplesGold.forEach { (max, inputs, program) ->
            assertThat(evaluateThrustersPowerWithFeedbackLoop(inputs, program)).isEqualTo(max)
        }
    }


    @Test
    fun findOptimum() {
        val max = listOf(0, 1, 2, 3, 4).permutations()
            .maxBy {
                evaluateThrustersPower(it, program)
            }!!

        println(evaluateThrustersPower(max, program))
    }

    @Test
    fun findOptimumGold() {
        val max = listOf(9, 8, 7, 6, 5).permutations()
            .maxBy {
                evaluateThrustersPowerWithFeedbackLoop(it, program)
            }!!

        println(evaluateThrustersPowerWithFeedbackLoop(max, program))
    }

    @Test
    fun testPermutations() {
        assertThat(listOf(0, 1, 2, 3, 4).permutations()).hasSize(120)
    }
}

fun <T> List<T>.permutations(): PersistentList<PersistentList<T>> {
    if (size == 1) return persistentListOf(toPersistentList())
    val droppedFirst = first()
    val permutationsOfTheRest = drop(1).permutations()

    return permutationsOfTheRest.flatMap { permutation ->
        // for every permutation move the dropped first through the list
        // droppedFirst = 1
        // drop(1).permutations().first = 2 3 4 5
        // 1 2 3 4 5
        // 2 1 3 4 5
        // 2 3 1 4 5
        // 2 3 4 1 5
        // 2 3 4 5 1
        // do the same for other permutations of 2 3 4 5
        (0..permutation.size).map { index ->
            permutation.add(index, droppedFirst)
        }
    }.toPersistentList()

}

private val examples: List<Triple<Int, List<Int>, PersistentList<Int>>> = listOf(
    Triple<Int, List<Int>, PersistentList<Int>>(
        43210,
        listOf<Int>(4, 3, 2, 1, 0),
        persistentListOf<Int>(3, 15, 3, 16, 1002, 16, 10, 16, 1, 16, 15, 15, 4, 15, 99, 0, 0)
    ),
    Triple<Int, List<Int>, PersistentList<Int>>(
        54321,
        listOf<Int>(0, 1, 2, 3, 4),
        persistentListOf<Int>(
            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
        )
    ),
    Triple<Int, List<Int>, PersistentList<Int>>(
        65210,
        listOf<Int>(1, 0, 4, 3, 2),
        persistentListOf<Int>(
            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
        )
    )
)

private val examplesGold: List<Triple<Int, List<Int>, PersistentList<Int>>> = listOf(
    Triple<Int, List<Int>, PersistentList<Int>>(
        139629729,
        listOf<Int>(9, 8, 7, 6, 5),
        persistentListOf<Int>(
            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
        )
    ),
    Triple<Int, List<Int>, PersistentList<Int>>(
        18216,
        listOf<Int>(9, 7, 8, 5, 6),
        persistentListOf<Int>(
            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
        )
    )
)

private val program = persistentListOf(
    3,
    8,
    1001,
    8,
    10,
    8,
    105,
    1,
    0,
    0,
    21,
    42,
    55,
    76,
    89,
    114,
    195,
    276,
    357,
    438,
    99999,
    3,
    9,
    1001,
    9,
    3,
    9,
    1002,
    9,
    3,
    9,
    1001,
    9,
    3,
    9,
    1002,
    9,
    2,
    9,
    4,
    9,
    99,
    3,
    9,
    102,
    2,
    9,
    9,
    101,
    5,
    9,
    9,
    4,
    9,
    99,
    3,
    9,
    102,
    3,
    9,
    9,
    101,
    5,
    9,
    9,
    1002,
    9,
    2,
    9,
    101,
    4,
    9,
    9,
    4,
    9,
    99,
    3,
    9,
    102,
    5,
    9,
    9,
    1001,
    9,
    3,
    9,
    4,
    9,
    99,
    3,
    9,
    1001,
    9,
    4,
    9,
    102,
    5,
    9,
    9,
    1001,
    9,
    5,
    9,
    1002,
    9,
    2,
    9,
    101,
    2,
    9,
    9,
    4,
    9,
    99,
    3,
    9,
    101,
    1,
    9,
    9,
    4,
    9,
    3,
    9,
    101,
    1,
    9,
    9,
    4,
    9,
    3,
    9,
    1001,
    9,
    1,
    9,
    4,
    9,
    3,
    9,
    1001,
    9,
    2,
    9,
    4,
    9,
    3,
    9,
    1002,
    9,
    2,
    9,
    4,
    9,
    3,
    9,
    101,
    1,
    9,
    9,
    4,
    9,
    3,
    9,
    1001,
    9,
    2,
    9,
    4,
    9,
    3,
    9,
    101,
    1,
    9,
    9,
    4,
    9,
    3,
    9,
    1002,
    9,
    2,
    9,
    4,
    9,
    3,
    9,
    1001,
    9,
    2,
    9,
    4,
    9,
    99,
    3,
    9,
    1001,
    9,
    2,
    9,
    4,
    9,
    3,
    9,
    101,
    2,
    9,
    9,
    4,
    9,
    3,
    9,
    1002,
    9,
    2,
    9,
    4,
    9,
    3,
    9,
    102,
    2,
    9,
    9,
    4,
    9,
    3,
    9,
    1002,
    9,
    2,
    9,
    4,
    9,
    3,
    9,
    102,
    2,
    9,
    9,
    4,
    9,
    3,
    9,
    102,
    2,
    9,
    9,
    4,
    9,
    3,
    9,
    101,
    1,
    9,
    9,
    4,
    9,
    3,
    9,
    101,
    1,
    9,
    9,
    4,
    9,
    3,
    9,
    1002,
    9,
    2,
    9,
    4,
    9,
    99,
    3,
    9,
    102,
    2,
    9,
    9,
    4,
    9,
    3,
    9,
    101,
    1,
    9,
    9,
    4,
    9,
    3,
    9,
    101,
    1,
    9,
    9,
    4,
    9,
    3,
    9,
    102,
    2,
    9,
    9,
    4,
    9,
    3,
    9,
    101,
    1,
    9,
    9,
    4,
    9,
    3,
    9,
    102,
    2,
    9,
    9,
    4,
    9,
    3,
    9,
    101,
    1,
    9,
    9,
    4,
    9,
    3,
    9,
    102,
    2,
    9,
    9,
    4,
    9,
    3,
    9,
    101,
    1,
    9,
    9,
    4,
    9,
    3,
    9,
    101,
    2,
    9,
    9,
    4,
    9,
    99,
    3,
    9,
    1002,
    9,
    2,
    9,
    4,
    9,
    3,
    9,
    1001,
    9,
    2,
    9,
    4,
    9,
    3,
    9,
    101,
    2,
    9,
    9,
    4,
    9,
    3,
    9,
    1001,
    9,
    1,
    9,
    4,
    9,
    3,
    9,
    101,
    2,
    9,
    9,
    4,
    9,
    3,
    9,
    101,
    1,
    9,
    9,
    4,
    9,
    3,
    9,
    1001,
    9,
    1,
    9,
    4,
    9,
    3,
    9,
    1001,
    9,
    2,
    9,
    4,
    9,
    3,
    9,
    102,
    2,
    9,
    9,
    4,
    9,
    3,
    9,
    1001,
    9,
    1,
    9,
    4,
    9,
    99,
    3,
    9,
    1001,
    9,
    1,
    9,
    4,
    9,
    3,
    9,
    101,
    1,
    9,
    9,
    4,
    9,
    3,
    9,
    1002,
    9,
    2,
    9,
    4,
    9,
    3,
    9,
    102,
    2,
    9,
    9,
    4,
    9,
    3,
    9,
    1002,
    9,
    2,
    9,
    4,
    9,
    3,
    9,
    101,
    2,
    9,
    9,
    4,
    9,
    3,
    9,
    1001,
    9,
    1,
    9,
    4,
    9,
    3,
    9,
    1002,
    9,
    2,
    9,
    4,
    9,
    3,
    9,
    102,
    2,
    9,
    9,
    4,
    9,
    3,
    9,
    101,
    2,
    9,
    9,
    4,
    9,
    99
)