package aoc2023

import println
import readInput

/*
--- Day 19: Aplenty ---
https://adventofcode.com/2023/day/19
*/
fun main() {

    val inlineTestInput = """
px{a<2006:qkq,m>2090:A,rfg}
pv{a>1716:R,A}
lnx{m>1548:A,A}
rfg{s<537:gd,x>2440:R,A}
qs{s>3448:A,lnx}
qkq{x<1416:A,crn}
crn{x>2662:A,R}
in{s<1351:px,qqz}
qqz{s>2770:qs,m<1801:hdj,R}
gd{a>3333:R,R}
hdj{m>838:A,pv}

{x=787,m=2655,a=1222,s=2876}
{x=1679,m=44,a=2067,s=496}
{x=2036,m=264,a=79,s=2244}
{x=2461,m=1339,a=466,s=291}
{x=2127,m=1623,a=2188,s=1013}
"""

    data class RuleCondition(val attr: Int = 0, val isSmaller: Boolean = false, val value: Int = 0, val follow: String)

    fun createMap(input: List<String>): Pair<Int, Map<String, List<RuleCondition>>> {
        val ruleMap = HashMap<String, List<RuleCondition>>()
        var p = 0
        do {
            "([a-z]+)\\{(.*)}".toRegex().matchEntire(input[p])?.destructured
                ?.let { (id, rules) ->
                    val rulesList = rules.split(",").map {
                        if (it.contains("<")) {
                            val (attr, rest) = it.split("<")
                            val (value, branchId) = rest.split(":")
                            RuleCondition("xmas".indexOf(attr), true, value.toInt(), branchId)
                        } else if (it.contains(">")) {
                            val (attr, rest) = it.split(">")
                            val (value, branchId) = rest.split(":")
                            RuleCondition("xmas".indexOf(attr), false, value.toInt(), branchId)
                        } else {
                            RuleCondition(follow = it)
                        }
                    }
                    ruleMap[id] = rulesList
                }
            if (input[p].isBlank()) break
            p++
        } while (true)
        return p to ruleMap
    }

    fun part1(input: List<String>): Int {
        val (p, ruleMap) = createMap(input)
        var sum = 0
        for (pp in p + 1 until input.size) {
            val attrs =
                input[pp].split("=", ",", "x", "m", "a", "s", "{", "}").filter { it.isNotBlank() }.map { it.toInt() }
                    .toIntArray()
            var ruleId = "in"
            while (ruleId != "A" && ruleId != "R") {
                val conditions = ruleMap[ruleId]!!
                for (cond in conditions) {
                    if (cond.isSmaller) {
                        if (attrs[cond.attr] < cond.value) {
                            ruleId = cond.follow
                            break
                        }
                    } else {
                        if (attrs[cond.attr] > cond.value) {
                            ruleId = cond.follow
                            break
                        }
                    }
                }
            }
            if (ruleId == "A") {
                sum += attrs.sum()
            }
        }
        return sum
    }

    fun recurse(ruleMap: Map<String, List<RuleCondition>>, ranges: List<IntRange>, ruleId: String): Long {
        if (ruleId == "A") {
            return ranges.map { (it.last - it.first + 1).coerceAtLeast(0).toLong() }.reduce(Long::times)
        }
        if (ruleId == "R") return 0
        var result = 0L
        val conditions = ruleMap[ruleId]!!
        var currRanges = ranges
        for (cond in conditions) {
            if (cond.value != 0) {
                if (cond.isSmaller) {
                    val subRange = currRanges.mapIndexed { index, intRange ->
                        if (index != cond.attr) intRange else
                            IntRange(intRange.first, intRange.last.coerceAtMost(cond.value - 1))
                    }
                    result += recurse(ruleMap, subRange, cond.follow)
                    currRanges = currRanges.mapIndexed { index, intRange ->
                        if (index != cond.attr) intRange else
                            IntRange(intRange.first.coerceAtLeast(cond.value), intRange.last)
                    }
                } else {
                    val subRange = currRanges.mapIndexed { index, intRange ->
                        if (index != cond.attr) intRange else
                            IntRange(intRange.first.coerceAtLeast(cond.value + 1), intRange.last)
                    }
                    result += recurse(ruleMap, subRange, cond.follow)
                    currRanges = currRanges.mapIndexed { index, intRange ->
                        if (index != cond.attr) intRange else
                            IntRange(intRange.first, intRange.last.coerceAtMost(cond.value))
                    }
                }
            } else {
                result += recurse(ruleMap, currRanges, cond.follow)
                break
            }
        }
        return result
    }

    fun part2(input: List<String>): Long {
        val (_, ruleMap) = createMap(input)
        val ranges = listOf(IntRange(1, 4000), IntRange(1, 4000), IntRange(1, 4000), IntRange(1, 4000))
        return recurse(ruleMap, ranges, "in")
    }

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

    val input = readInput("aoc2023/Day19")
    part1(input).println()
    part2(input).println()
}
