package be.vandewalleh.aoc.days

import be.vandewalleh.aoc.utils.input.Day
import be.vandewalleh.aoc.utils.input.Input
import be.vandewalleh.aoc.utils.input.Lines
import be.vandewalleh.aoc.utils.input.createDay
import kotlin.math.pow
import org.eclipse.collections.impl.factory.primitive.IntObjectMaps
import org.eclipse.collections.impl.factory.primitive.LongIntMaps

@Day(14)
class Day14(@Lines val input: Input<List<String>>) {

    private val memRe = "mem\\[(\\d+)] = (.*)$".toRegex()

    private fun Long.toBin36() = toString(2).padStart(length = 36, padChar = '0')

    fun part1(): Long {
        val mem = IntObjectMaps.mutable.empty<String>()

        var currentMask: String = ""

        for (line in input.value) {
            if (line.startsWith("mask")) {
                currentMask = line.removePrefix("mask = ")
            } else {
                val (address, value) = memRe.find(line)!!.destructured
                val bin = value.toLong().toBin36()

                val result = bin.zip(currentMask)
                    .map { (bin, mask) -> if (mask != 'X') mask else bin }
                    .joinToString("")

                mem.put(address.toInt(), result)
            }
        }

        return mem.values()
            .map { it.dropWhile { it == '0' } }
            .map { it.toLong(2) }
            .sum()
    }

    fun part2(): Long {
        val mem = LongIntMaps.mutable.empty()

        var currentMask = ""

        for (line in input.value) {
            if (line[1] == 'a') {
                currentMask = line.substring(7)
            } else {

                val (address, value) = memRe.find(line)!!.destructured.let { (add, value) ->
                    add.toLong().toBin36() to value.toInt()
                }

                val mutations = generateMutations(currentMask, address)

                for (mutation in mutations) {
                    mem.put(String(mutation).toLong(2), value)
                }
            }
        }

        return mem.values().sum()
    }

    private fun generateMutations(mask: String, address: String): Array<CharArray> {
        val mutationCount = mask.count { it == 'X' }.let { 2.0.pow(it.toDouble()).toInt() }
        val mutations = Array(mutationCount) { CharArray(36) }

        var groups = 1

        for (i in mask.indices) {
            when (mask[i]) {
                'X' -> {
                    groups *= 2
                    val groupSize = mutationCount / groups
                    var currentChar = '1'
                    for (b in mutations.indices) {
                        val builder = mutations[b]
                        val flip = b % groupSize == 0
                        if (flip) currentChar = if (currentChar == '0') '1' else '0'
                        builder[i] = currentChar
                    }
                }
                '1' -> mutations.forEach { it[i] = '1' }
                else -> mutations.forEach { it[i] = address[i] }
            }
        }
        return mutations
    }

}

fun main() = with(createDay<Day14>()) {
    println(part1())
    println(part2())
}