2021 Day 11

class Octopus(var energyLevel: Int) {
    var surroundingOctopuses: List<Octopus> = emptyList()

    fun increaseEnergy() {
        energyLevel++
        if (energyLevel == 10)
            surroundingOctopuses.forEach { it.increaseEnergy() }
    }

    fun flash(): Int {
        if (energyLevel < 10 ) return 0
        energyLevel = 0
        return 1
    }
}

class OctopusCave {
    @Test
    fun `Should increase energy one at a time`() {
        val octopus = Octopus(1)
        octopus.increaseEnergy()
        assertEquals(2, octopus.energyLevel)
        octopus.increaseEnergy()
        assertEquals(3, octopus.energyLevel)
    }

    @Test
    fun `Should drop to zero on flash`() {
        val octopus = Octopus(10)
        octopus.flash()
        assertEquals(0, octopus.energyLevel)
    }

    @Test
    fun `Should reset to 0 if flashing above 9 energy`() {
        val octopus = Octopus(7)
        octopus.flash()
        assertEquals(7, octopus.energyLevel)
        octopus.increaseEnergy()
        octopus.flash()
        assertEquals(8, octopus.energyLevel)
        octopus.increaseEnergy()
        octopus.flash()
        assertEquals(9, octopus.energyLevel)
        octopus.increaseEnergy()
        octopus.flash()
        assertEquals(0, octopus.energyLevel)
    }

    @Test
    fun `Should increase adjacent energy when increasing above 9`() {
        val octopus = Octopus(9)
        val neighbor1 = Octopus(5)
        val neighbor2 = Octopus(9)
        val neighbor3 = Octopus(4)
        val someOther = Octopus(7)
        octopus.surroundingOctopuses = listOf(neighbor1, neighbor2, neighbor3)
        neighbor1.surroundingOctopuses = listOf(octopus, neighbor2, someOther)
        neighbor2.surroundingOctopuses = listOf(octopus, neighbor1)
        octopus.increaseEnergy()
        assertEquals(11, octopus.energyLevel)
        assertEquals(7, neighbor1.energyLevel)
        assertEquals( 10, neighbor2.energyLevel)
        assertEquals( 5, neighbor3.energyLevel)
        assertEquals( 7, someOther.energyLevel)
    }

    @Test
    fun `Small example single step` () {
        assertEquals(9, smallExampleCave.step())
        assertEquals(0, smallExampleCave.step())
    }

    @Test
    fun `Large example 10 steps` () {
        var totalFlashes = 0L
        repeat(10) { totalFlashes += largeExampleCave.step() }
        assertEquals(204, totalFlashes)
    }

    @Test
    fun `Large example 100 steps` () {
        var totalFlashes = 0L
        repeat(100) { totalFlashes += largeExampleCave.step() }
        assertEquals(1656, totalFlashes)
    }

    @Test
    fun `Part 1` () {
        var totalFlashes = 0L
        repeat(100) { totalFlashes += puzzleCave.step() }
        println("Part 1: $totalFlashes")
    }

    @Test
    fun `First synchronized flash` () {
        var step = 1
        while (largeExampleCave.step() != largeExampleCave.size) step++
        assertEquals(195, step)
    }

    @Test
    fun `Part 2` () {
        var step = 1
        while (puzzleCave.step() != puzzleCave.size) step++
        println("Part 2: $step")
    }

    private fun List<List<Octopus>>.withSurrounding(): List<Octopus> {
        this.forEachIndexed { y, row ->
            row.forEachIndexed{ x, octopus ->
                octopus.surroundingOctopuses = surrounding(x, y)
            }
        }
        return this.flatten();
    }

    private fun List<List<Octopus>>.surrounding(x: Int, y: Int): List<Octopus> {
        var list = emptyList<Octopus>().toMutableList()
        // Three octopuses above
        if (y != 0) {
            if (x != 0) list.add(this[y-1][x-1])
            list.add(this[y-1][x])
            if (x != this[0].size-1) list.add(this[y-1][x+1])
        }
        // Two octopuses left and right
        if (x != 0) list.add(this[y][x-1])
        if (x != this[0].size-1) list.add(this[y][x+1])
        // Three octopuses below
        if (y != this.size-1) {
            if (x != 0) list.add(this[y+1][x-1])
            list.add(this[y+1][x])
            if (x != this[0].size-1) list.add(this[y+1][x+1])
        }
        return list
    }

    private fun List<Octopus>.step(): Int {
        this.forEach { it.increaseEnergy() }
        return this.sumBy { it.flash() }
    }

    private val smallExampleCave = """11111
19991
19191
19991
11111""".split("\n").map { it.map { it.toString().toInt() }.map { Octopus(it) } }.withSurrounding()

    private val largeExampleCave = """5483143223
2745854711
5264556173
6141336146
6357385478
4167524645
2176841721
6882881134
4846848554
5283751526""".split("\n").map { it.map { it.toString().toInt() }.map { Octopus(it) } }.withSurrounding()

    private val puzzleCave = """5723573158
3154748563
4783514878
3848142375
3637724151
8583172484
7747444184
1613367882
6228614227
4732225334""".split("\n").map { it.map { it.toString().toInt() }.map { Octopus(it) } }.withSurrounding()
}