package aoc2024

import println
import readInput
import java.util.*

/*
--- Day 9: Disk Fragmenter ---
https://adventofcode.com/2024/day/9
*/
fun main() {

    val inlineTestInput = """2333133121414131402"""

    fun part1(input: List<String>): Long {
        val diskSize = input[0].sumOf { it - '0' }
        val bitmap = IntArray(diskSize) { -1 }
        var pos = 0
        var isFile = true
        var fileId = 0
        var usedBlocks = 0
        for (c in input[0]) {
            val size = c - '0'
            if (isFile) {
                usedBlocks += size
                for (p in 1..size) {
                    bitmap[pos++] = fileId
                }
                fileId++
            } else {
                pos += size
            }
            isFile = !isFile
        }
        var backPos = diskSize - 1
        for (i in 0 until usedBlocks) {
            if (bitmap[i] == -1) {
                while (bitmap[backPos] < 0) {
                    backPos--
                }
                bitmap[i] = bitmap[backPos]
                bitmap[backPos--] = -1
            }
        }
        return bitmap.take(usedBlocks).map(Int::toLong).reduceIndexed { index, acc, i -> acc + index * i }
    }

    fun part2(input: List<String>): Long {
        var pos = 0
        var isFile = true
        val freeList = Array<MutableSet<Int>>(10) { TreeSet() }
        val fileList = ArrayList<Pair<Int, Int>>((input[0].length + 1) / 2)
        for (c in input[0]) {
            val size = c - '0'
            if (isFile) {
                fileList.add(pos to size)
            } else if (size > 0) {
                freeList[size].add(pos)
            }
            pos += size
            isFile = !isFile
        }

        val newFileList = Array<Pair<Int, Int>?>(fileList.size) { null }
        var revPos = fileList.size
        for (file in fileList.asReversed()) {
            val (filePos, size) = file
            var bestListSize = 0
            var bestListPos = 0
            for (s in size..9) {
                val list = freeList[s]
                val firstPos = list.firstOrNull()
                if ((firstPos != null && firstPos < filePos) &&
                    (bestListSize == 0 || firstPos < bestListPos)
                ) {
                    bestListSize = s
                    bestListPos = firstPos
                }
            }
            if (bestListSize != 0) {
                val newPos = freeList[bestListSize].first()
                freeList[bestListSize].remove(newPos)
                if (bestListSize > size) freeList[bestListSize - size].add(newPos + size)
                newFileList[--revPos] = newPos to size
            } else {
                newFileList[--revPos] = file
                // stop if not even a size 1 file could be fit
                if (size == 1) break
            }
        }
        while (revPos >= 0) {
            newFileList[revPos] = fileList[revPos]
            revPos--
        }
        // Gauß to the rescue!
        return newFileList.mapIndexed { index, file -> index * (file!!.first.toLong() * file.second + (file.second * (file.second - 1) / 2)) }
            .sum()
    }

    // test if implementation meets criteria from the description, like:
    val testInput = inlineTestInput.trim().reader().readLines()
    //val testInput = readInput("aoc2024/Day09_test")
    val testInputPart1Result = part1(testInput)
    println("Part 1 Test: $testInputPart1Result")
    val testInputPart2Result = part2(testInput)
    println("Part 2 Test: $testInputPart2Result")
    check(testInputPart1Result == 1928L)
    check(testInputPart2Result == 2858L)

    val input = readInput("aoc2024/Day09")
    part1(input).println()
    part2(input).println()
}
