import java.io.File

abstract class Tree(val weight: Int) {
    abstract fun walk(parent: Tree?, name: String, function: (Tree?, String, Tree) -> Unit)
    abstract fun <T> reduce(transformFunction: (Tree) -> T, combineFunction: (T, List<T>) -> T): T

    val combinedWeight: Int
        get() {
            val reduceFunction: (Tree) -> Int = {
                it.weight
            }
            val combineFunction: (Int, List<Int>) -> Int = {
                parent, children ->
                parent + children.fold(0) { a, b -> a + b}
            }
            return reduce(reduceFunction, combineFunction)
        }
}
class Placeholder : Tree(0) {
    override fun <T> reduce(transformFunction: (Tree) -> T, combineFunction: (T, List<T>) -> T): T {
        return transformFunction(this)
    }

    override fun walk(parent: Tree?, name: String, function: (Tree?, String, Tree) -> Unit) {
        function(parent, name, this)
    }
}

class Child(weight: Int) : Tree(weight) {
    override fun <T> reduce(transformFunction: (Tree) -> T, combineFunction: (T, List<T>) -> T): T {
        return transformFunction(this)
    }

    override fun walk(parent: Tree?, name: String, function: (Tree?, String, Tree) -> Unit) {
        function(parent, name, this)
    }
}

class Parent(weight: Int) : Tree(weight) {
    val children = mutableMapOf<String, Tree>()

    override fun <T> reduce(transformFunction: (Tree) -> T, combineFunction: (T, List<T>) -> T): T {
        return combineFunction(transformFunction(this), children.map { it.value.reduce(transformFunction, combineFunction) })
    }

    override fun walk(parent: Tree?, name: String, function: (Tree?, String, Tree) -> Unit) {
        function(parent, name, this)
        children.forEach { t, u -> u.walk(this, t, function) }
    }
}

fun main(args: Array<String>) {
    val extractionPattern = Regex("([a-z]+)\\s*\\((\\d+)\\)\\s*(?:->\\s*(.+))?")
    val baseMap = mutableMapOf<String, Tree>()
    val nodesWithParent = mutableListOf<Tree>()
    File("../puzzle_7.txt").readLines().forEach {
        val match = extractionPattern.matchEntire(it)!!
        val name = match.groupValues[1]
        val weight = match.groupValues[2].toInt()
        val childrenString = match.groupValues[3]
        val children = if(childrenString.isEmpty()) emptyList() else childrenString.split(", ")
        val newNode = if(children.isEmpty()) {
            Child(weight)
        } else {
            val newParent = Parent(weight)
            children.forEach { childName -> newParent.children.put(childName, Placeholder()) }
            newParent
        }
        baseMap.put(name, newNode)
    }
    baseMap.forEach {
        val treeNode = it.value as? Parent ?: return@forEach

        treeNode.children.forEach {
            pair ->
            val (childKey, _) = pair
            val childNode = baseMap[childKey]!!
            treeNode.children[childKey] = childNode
            nodesWithParent.add(childNode)
        }
    }
    val (name, node) = baseMap.filterValues { !nodesWithParent.contains(it) }.entries.first()
    node.walk(null, name) {
        parent, name, node ->
        if(parent == null) println("root node is $name")
        if(node is Parent){
            val childWeightPairs = node.children.map { it.value.combinedWeight to it.key }
            val childWeights = childWeightPairs.map { it.first }
            if(childWeights.toSet().size != 1) {
                val frequencyMap = mutableMapOf<Int, Int>()
                childWeights.forEach { frequencyMap[it] = frequencyMap.getOrDefault(it, 0) + 1 }
                val culpritWeight = frequencyMap.filter { it.value == 1 }.keys.first()
                val properWeight = frequencyMap.filter { it.key != culpritWeight }.keys.first()
                val culpritNode = childWeightPairs.filter { it.first == culpritWeight }.first()
                val culpritChildren = (node.children[culpritNode.second] as Parent).children.values
                val culpritChildrenWeights = culpritChildren.map { it.combinedWeight }
                if(culpritChildrenWeights.toSet().size == 1) {
                    println("$culpritNode has balanced children, but doesn't weight right.")
                    println("Correct weight is $properWeight, but actually is $culpritWeight")
                    val difference = properWeight - culpritWeight
                    println("The difference is $difference, and must be adjusted with node weight.")
                    println("Correct node weight is ${difference + node.children[culpritNode.second]!!.weight}")
                }
            }
        }
    }
}