package aoc2023

import CharGrid
import RelPos
import println
import readInput

/*
--- Day 21: Step Counter ---
https://adventofcode.com/2023/day/21
*/
fun main() {

    val inlineTestInput = """
...........
.....###.#.
.###.##..#.
..#.#...#..
....#.#....
.##..S####.
.##..#...#.
.......##..
.##.#.####.
.##..##.##.
...........
"""

    fun part1(input: List<String>, steps: Int): Int {
        var grid = CharGrid(input, '#')
        for (i in 1..steps) {
            grid = grid.grow(CharGrid.PLUS_POS, { it == 'S' }, { org: Char, old: Char, _: RelPos ->
                if (old != '#') 'O' else old
            })
            grid.apply {
                when (it) {
                    'S' -> '.'
                    'O' -> 'S'
                    else -> it
                }
            }
        }
        return grid.findMatches { it == 'S' }.count()
    }

    fun fn(x: Long) = 61124L * x * x - 91368L * x + 34140L

    fun part2(input: List<String>, steps: Int): Long {
        val grid = CharGrid(input, '%')
        var exp = grid.findMatches { it == 'S' }.map { RelPos(it.first, it.second) }.toSet()
        val mp = steps % 262
        for (i in 1..(262 * 2 + mp)) {
            exp =
                exp.flatMap {
                    CharGrid.PLUS_POS.map { rp -> it.translate(rp) }
                        .filter { grid[Math.floorMod(it.dr, grid.width), Math.floorMod(it.dc, grid.height)] != '#' }
                }.toSet()
            val count = exp.size
            if (i % 262 == mp) println("$i: $count")
        }
        println(fn(1L))
        println(fn(2L))
        println(fn(3L))

        // Wolfram alpha says: {3896, 95900, 310152} -> 61124 x^2 - 91368 x + 34140
        return fn((steps / 262L) + 1)
    }

    // test if implementation meets criteria from the description, like:
    val testInput = inlineTestInput.trim().reader().readLines()
    //val testInput = readInput("aoc2023/Day21_test")
    val testInputPart1Result = part1(testInput, 6)
    println("Part 1 Test: $testInputPart1Result")
    //val testInputPart2Result = part2(testInput, 500)
    //println("Part 2 Test: $testInputPart2Result")
    check(testInputPart1Result == 16)
    //check(testInputPart2Result == 167004L)

    val input = readInput("aoc2023/Day21")
    part1(input, 64).println()
    part2(input, 26501365).println()
}
