package aoc2023

import lcm
import println
import readInput

/*
--- Day 20: Pulse Propagation ---
https://adventofcode.com/2023/day/20
*/
fun main() {

    val inlineTestInput = """
broadcaster -> a, b, c
%a -> b
%b -> c
%c -> inv
&inv -> a
"""

    val inlineTestInput2 = """
broadcaster -> a
%a -> inv, con
&inv -> b
%b -> con
&con -> output
"""

    data class Node(
        val id: String,
        val targets: List<String>,
        val type: Int,
        val inputs: MutableList<String> = ArrayList(),
        var signalsIn: MutableList<Boolean> = ArrayList(),
        var stateIdx: Int = 0
    )

    data class CacheInfo(val lowPulses: Long, val highPulses: Long, val nextState: BooleanArray)

    fun createGraph(input: List<String>): HashMap<String, Node> {
        val nodes = HashMap<String, Node>()
        for (i in input) {
            val (source, targets) = i.split(" -> ")
            val targetList = targets.split(", ")
            val (id, type) = if (source.startsWith("%")) {
                source.removePrefix("%") to 1
            } else if (source.startsWith("&")) {
                source.removePrefix("&") to 2
            } else {
                source to 0
            }
            nodes[id] = Node(id, targetList, type)
        }
        var stateIdx = 0
        for (node in nodes.values) {
            node.stateIdx = stateIdx++
            for (t in node.targets) {
                nodes[t]?.inputs?.add(node.id)
            }
        }
        return nodes
    }

    fun part1(input: List<String>): Long {
        val nodes = createGraph(input)
        val state = BooleanArray(nodes.size)
        val cache = HashMap<String, CacheInfo>()
        var totalLowPulses = 0L
        var totalHighPulses = 0L
        for (n in 1..1000) {
            val stateString = state.joinToString("") { if (it) "1" else "0" }
            val cached = cache[stateString]
            if (cached != null) {
                totalLowPulses += cached.lowPulses
                totalHighPulses += cached.highPulses
                cached.nextState.copyInto(state)
                continue
            }
            val queue = LinkedHashSet<String>()
            queue.add("broadcaster")
            nodes["broadcaster"]!!.signalsIn.add(false)
            var lowPulses = 1L
            var highPulses = 0L
            do {
                val id = queue.first()
                queue.remove(id)
                val node = nodes[id]
                if (node != null) {
                    val signalsIn = node.signalsIn
                    for (signalIn in signalsIn) {
                        when (node.type) {
                            0 -> {
                                state[node.stateIdx] = signalIn
                                if (state[node.stateIdx]) highPulses += node.targets.size else lowPulses += node.targets.size
                                node.targets.forEach { nodes[it]?.signalsIn?.add(state[node.stateIdx]) }
                                queue.addAll(node.targets)
                            }

                            1 -> if (!signalIn) {
                                state[node.stateIdx] = !state[node.stateIdx]
                                if (state[node.stateIdx]) highPulses += node.targets.size else lowPulses += node.targets.size
                                node.targets.forEach { nodes[it]?.signalsIn?.add(state[node.stateIdx]) }
                                queue.addAll(node.targets)
                            }

                            2 -> {
                                state[node.stateIdx] = !node.inputs.all { state[nodes[it]!!.stateIdx] }
                                if (state[node.stateIdx]) highPulses += node.targets.size else lowPulses += node.targets.size
                                node.targets.forEach { nodes[it]?.signalsIn?.add(state[node.stateIdx]) }
                                queue.addAll(node.targets)
                            }
                        }
                    }
                    node.signalsIn.clear()
                }
            } while (queue.isNotEmpty())
            cache[stateString] = CacheInfo(lowPulses, highPulses, state.copyOf())
            totalLowPulses += lowPulses
            totalHighPulses += highPulses
        }
        return totalLowPulses * totalHighPulses
    }

    fun part2(input: List<String>): Long {
        val nodes = createGraph(input)
        val state = BooleanArray(nodes.size)
        var buttonCount = 0L
        val klNode = nodes.values.first { it.targets.contains("rx") }
        val periods = LongArray(klNode.inputs.size)
        do {
            buttonCount++
            val queue = LinkedHashSet<String>()
            queue.add("broadcaster")
            nodes["broadcaster"]!!.signalsIn.add(false)
            do {
                val id = queue.first()
                queue.remove(id)
                val node = nodes[id]
                if (node != null) {
                    val signalsIn = node.signalsIn
                    for (signalIn in signalsIn) {
                        when (node.type) {
                            0 -> {
                                state[node.stateIdx] = signalIn
                                node.targets.forEach { nodes[it]?.signalsIn?.add(state[node.stateIdx]) }
                                queue.addAll(node.targets)
                            }

                            1 -> if (!signalIn) {
                                state[node.stateIdx] = !state[node.stateIdx]
                                node.targets.forEach { nodes[it]?.signalsIn?.add(state[node.stateIdx]) }
                                queue.addAll(node.targets)
                            }

                            2 -> {
                                state[node.stateIdx] = !node.inputs.all { state[nodes[it]!!.stateIdx] }
                                node.targets.forEach { nodes[it]?.signalsIn?.add(state[node.stateIdx]) }
                                queue.addAll(node.targets)
                            }
                        }
                    }
                    node.signalsIn.clear()
                    klNode.inputs.forEachIndexed { index, s ->
                        if (state[nodes[s]!!.stateIdx]) {
                            if (periods[index] == 0L) {
                                periods[index] = buttonCount
                                println("$index $s $buttonCount")
                            }
                        }
                    }

                    //node.targets.forEach { println("$rxCount: $lowPulses, $highPulses   ${node.id} -${state[node.stateIdx]}> $it") }
                }
            } while (queue.isNotEmpty())
            //println("$buttonCount")
        } while (periods.any { it == 0L })
        return periods.asIterable().lcm()
    }

    // test if implementation meets criteria from the description, like:
    val testInput = inlineTestInput.trim().reader().readLines()
    val testInput2 = inlineTestInput2.trim().reader().readLines()
    //val testInput = readInput("aoc2023/Day20_test")
    val testInputPart1Result = part1(testInput)
    println("Part 1 Test: $testInputPart1Result")
    val testInputPart1bResult = part1(testInput2)
    println("Part 1b Test: $testInputPart1bResult")
    check(testInputPart1Result == 32000000L)
    check(testInputPart1bResult == 11687500L)

    val input = readInput("aoc2023/Day20")
    part1(input).println()
    part2(input).println()
}
