package be.vandewalleh.aoc.days

import be.vandewalleh.aoc.utils.input.Day
import be.vandewalleh.aoc.utils.input.Lines

private data class Point(val x: Int, val y: Int, val z: Int)
private data class Point4(val x: Int, val y: Int, val z: Int, val blah: Int)

private enum class State { Active, Inactive }

@Day(17)
class Day17(@Lines val input: List<String>) {

    fun part1(): Int {
        val grid = parseGrid { x, y -> Point(x, y, 0) }
        repeat(6) { step(grid, ::neighbours3) }
        return grid.values.count { it == State.Active }
    }

    fun part2(): Int {
        val grid = parseGrid { x, y -> Point4(x, y, 0, 0) }
        repeat(6) { step(grid, ::neighbours4) }
        return grid.values.count { it == State.Active }
    }

    private fun <T> parseGrid(pointFactory: (x: Int, y: Int) -> T): MutableMap<T, State> {
        val grid = mutableMapOf<T, State>()
        input.forEachIndexed { index, row ->
            row.forEachIndexed { col, char ->
                val state = if (char == '#') State.Active else State.Inactive
                grid[pointFactory(index, col)] = state
            }
        }
        return grid
    }

    private fun <T> step(grid: MutableMap<T, State>, neighbours: (T) -> List<T>) {
        val modifications = mutableMapOf<T, State>()
        val pointsToConsider = grid.keys.flatMap { neighbours(it) }.toSet()
        for (point in pointsToConsider) {
            val neighbours = neighbours(point)
            val activeNeighboursCount = neighbours.count { grid[it] ?: State.Inactive == State.Active }
            val state = grid[point] ?: State.Inactive
            if (state == State.Active && activeNeighboursCount !in 2..3) {
                modifications[point] = State.Inactive
            } else if (activeNeighboursCount == 3) {
                modifications[point] = State.Active
            }
        }
        for ((point, state) in modifications) {
            grid[point] = state
        }
    }

    private fun neighbours3(point: Point): List<Point> {
        val points = mutableListOf<Point>()
        for (x in point.x - 1..point.x + 1) {
            for (y in point.y - 1..point.y + 1) {
                for (z in point.z - 1..point.z + 1) {
                    val generatedPoint = Point(x, y, z)
                    if (generatedPoint != point)
                        points.add(generatedPoint)
                }
            }
        }
        check(points.size == 26) { "Points size was ${points.size}" }
        return points
    }

    private fun neighbours4(point: Point4): List<Point4> {
        val points = mutableListOf<Point4>()
        for (x in point.x - 1..point.x + 1) {
            for (y in point.y - 1..point.y + 1) {
                for (z in point.z - 1..point.z + 1) {
                    for (blah in point.blah - 1..point.blah + 1) {
                        val generatedPoint = Point4(x, y, z, blah)
                        if (generatedPoint != point)
                            points.add(generatedPoint)
                    }
                }
            }
        }
        check(points.size == 80) { "Points size was ${points.size}" }
        return points
    }

}