package be.vandewalleh.aoc.days

import be.vandewalleh.aoc.days.geometry.Grid
import be.vandewalleh.aoc.days.geometry.gridOf
import be.vandewalleh.aoc.days.geometry.transformations
import be.vandewalleh.aoc.utils.input.Day
import be.vandewalleh.aoc.utils.input.Groups
import kotlin.math.sqrt

private typealias Tile = Grid<Char>

@Day(20)
class Day20(@Groups val input: List<List<String>>) {
    private val tiles: Map<Int, Tile> = input
        .map { it[0].let { it.substring(5 until it.indexOf(':')).toInt() } to it.drop(1) }
        .associate { (id, tile) -> id to gridOf(tile) }

    private fun Tile.allEdges() = listOf(edges(), edges().map { it.reversed() }).flatten()

    private fun edgesMatch(a: Tile, b: Tile): Boolean {
        val edges = b.allEdges()
        return a.allEdges().any { a -> edges.any { a == it } }
    }

    private fun combine(tiles: List<Tile>) = sequence {
        for (i in 0 until tiles.lastIndex) {
            val a = tiles[i]
            for (j in i + 1 until tiles.size) {
                val b = tiles[j]
                yield(a to b)
            }
        }
    }

    private fun tilesByNeighbourCount(): MutableMap<Int, MutableList<Tile>> {
        val neighbours = combine(tiles.values.toList())
            .filter { (a, b) -> edgesMatch(a, b) }
            .map { it.toList() }
            .flatten()
            .groupBy { it }
            .values

        val map = mutableMapOf<Int, MutableList<Tile>>()
        for (neighbour in neighbours) {
            map.computeIfAbsent(neighbour.size) { mutableListOf() }.add(neighbour.first())
        }
        return map
    }

    fun part1() = tilesByNeighbourCount()[2]!!
        .map { tile -> tiles.entries.find { it.value == tile }!!.key.toLong() }
        .reduce { acc, id -> acc * id }

    private fun neighboursCount(grid: Grid<*>, x: Int, y: Int): Int {
        var neighboursCount = 2
        if (x != 0 && x != grid.lastColumnIndex) neighboursCount++
        if (y != 0 && y != grid.lastRowIndex) neighboursCount++
        return neighboursCount
    }

    private fun isLeftOk(grid: Grid<Tile?>, x: Int, y: Int, candidate: Tile) = grid[x - 1, y]
        ?.let { left -> candidate.firstColumn() == left.lastColumn() }
        ?: true

    private fun isTopOk(grid: Grid<Tile?>, x: Int, y: Int, candidate: Tile) = grid[x, y - 1]
        ?.let { top -> candidate.firstRow() == top.lastRow() }
        ?: true

    private fun isBottomOk(grid: Grid<*>, x: Int, y: Int, candidate: Tile, neighbours: Map<Int, List<Tile>>) =
        if (y == grid.lastRowIndex) true
        else neighbours[neighboursCount(grid, x, y + 1)]!!
            .filterNot { it == candidate }
            .count { it.transformations().any { candidate.lastColumn() == it.firstColumn() } } == 1

    private fun isRightOk(grid: Grid<*>, x: Int, y: Int, candidate: Tile, neighbours: Map<Int, List<Tile>>) =
        if (x == grid.lastColumnIndex) true
        else neighbours[neighboursCount(grid, x + 1, y)]!!
            .filterNot { it == candidate }
            .count { it.transformations().any { candidate.lastRow() == it.firstRow() } } == 1

    private fun assembleGrid(): Grid<Tile?> {
        val size = sqrt(tiles.size.toDouble()).toInt()
        val grid: Grid<Tile?> = Grid(Array(size) { Array(size) { null } })

        val neighbours: MutableMap<Int, MutableList<Tile>> = tilesByNeighbourCount()

        for (y in 0 until grid.height) {
            for (x in 0 until grid.width) {
                val neighboursCount = neighboursCount(grid, x, y)
                val candidates = neighbours[neighboursCount]!!

                val conditions = mutableListOf<(Tile) -> Boolean>()
                conditions += { isLeftOk(grid, x, y, it) }
                conditions += { isTopOk(grid, x, y, it) }

                // why is this condition needed ?
                if (x == 0 && y == 0) {
                    conditions += { isBottomOk(grid, x, y, it, neighbours) }
                    conditions += { isRightOk(grid, x, y, it, neighbours) }
                }

                var found: Tile? = null

                outer@ for (candidate in candidates) {
                    for (transform in candidate.transformations()) {
                        if (conditions.all { it(transform) }) {
                            found = transform
                            candidates.remove(candidate)
                            break@outer
                        }
                    }
                }

                check(found != null)
                grid[x, y] = found
            }
        }

        return grid
    }

    private fun removeGaps(grid: Grid<Tile?>) {
        for (y in 0 until grid.height) {
            for (x in 0 until grid.width) {
                grid[x, y] = removeGaps(grid[x, y]!!)
            }
        }
    }

    private fun removeGaps(tile: Tile): Tile {
        val oldData: ArrayList<ArrayList<Char>> = tile.data
        val newData = ArrayList<ArrayList<Char>>(oldData.size - 2)
        oldData.subList(1, oldData.size - 1).forEach { d ->
            val l = ArrayList<Char>().apply {
                addAll(d.subList(1, d.size - 1))
            }
            newData.add(l)
        }
        return Tile(newData)
    }

    private fun gridToTile(grid: Grid<Tile?>): Tile {
        val newData = ArrayList<ArrayList<Char>>()
        for (y in 0 until grid.height) {
            val row = grid.row(y)
            for (yy in 0 until row[0]!!.height) {
                val combinedRow = ArrayList<Char>().also { newData.add(it) }
                row.forEach { combinedRow.addAll(it!!.row(yy)) }
            }
        }
        return Tile(newData)
    }

    private fun Tile.subGridData(startX: Int, startY: Int, width: Int, height: Int): List<List<Char>> {
        val newData = ArrayList<ArrayList<Char>>()
        for (y in startY until startY + height) {
            val row = ArrayList<Char>().also { newData.add(it) }
            for (x in startX until startX + width) {
                row.add(this[x, y]!!)
            }
        }
        return newData
    }

    private fun List<List<Char>>.isMonster(monster: List<List<Char>>): Boolean {
        val monsterRe = monster.joinToString("") { it.joinToString("") }.replace(" ", ".").toRegex()
        val str = this.joinToString("") { it.joinToString("") }
        return monsterRe.matches(str)
    }

    fun part2(): Int {
        val grid = assembleGrid()
        removeGaps(grid)
        val megaTile = gridToTile(grid)

        val monster: List<List<Char>> = """
            |                  # |
            |#    ##    ##    ###|
            | #  #  #  #  #  #   |
        """.trimMargin("|").lines().map { it.toCharArray().dropLast(1) }

        val monsterWidth = monster[0].size
        val monsterHeight = 3
        val monsterSquares = monster.flatten().count { it == '#' }
        val squares = megaTile.data.flatten().count { it == '#' }

        for (g in megaTile.transformations()) {
            var count = 0
            for (y in 0 until g.lastRowIndex - monsterHeight) {
                for (x in 0 until g.lastColumnIndex - monsterWidth) {
                    val subgrid = g.subGridData(x, y, monsterWidth, monsterHeight)
                    if (subgrid.isMonster(monster)) {
                        count++
                    }
                }
            }
            if (count != 0) return squares - (count * monsterSquares)
        }
        error("No monsters found")
    }

}