import kotlin.collections.any
import kotlin.math.pow

fun main() {
    fun part1(input: List<String>): Long {
        val operations = setOf(CalibrationOperation.Plus, CalibrationOperation.Multiply)
        return generalizedSolution(input, operations)
    }

    fun part2(input: List<String>): Long {
        val operations = setOf(CalibrationOperation.Plus, CalibrationOperation.Multiply, CalibrationOperation.Concat)
        return generalizedSolution(input, operations)
    }

    val testInput = readInput("Day07_test")
    check(part1(testInput) == 3749L)
    check(part2(testInput) == 11387L)

    val input = readInput("Day07")
    part1(input).println()
    part2(input).println()
}

fun parseInputDay7(input: List<String>) = input.map {
    val calibrationResultAndInput = it.split(':')
    calibrationResultAndInput[0].toLong() to calibrationResultAndInput[1].split(' ').filter { it != "" }.map { it.toLong() }
}

fun generalizedSolution(input: List<String>, operations: Set<CalibrationOperation>): Long {
    val parsedInput = parseInputDay7(input)
    val operationsPermutations = CalibrationOperation.operationPermutationSequence(*operations.toTypedArray()).take(calculatePermutationsNeeded(parsedInput, operations)).toList()
    return sumOfPossibleCalibrationEquations(parsedInput, operationsPermutations)
}

fun calculatePermutationsNeeded(parsedInput: List<Pair<Long, List<Long>>>, operations: Set<CalibrationOperation>): Int {
    val highestNumberOfOperations = parsedInput.maxOf { it.second.size - 1 }
    return (1..highestNumberOfOperations).sumOf { operations.size.toDouble().pow(it).toInt() }
}

fun sumOfPossibleCalibrationEquations(parsedInput: List<Pair<Long, List<Long>>>, operationPermutationCollection: Collection<OperationPermutation>): Long {
    val permutationsGrouped = operationPermutationCollection.groupBy { it.size }
    return parsedInput.sumOf { (equationResult, equationInput) ->
        if (permutationsGrouped[equationInput.size - 1]!!.any { operations ->
                equationResult == equationInput.drop(1)
                    .foldIndexed(equationInput[0]) { index, acc, lng -> operations[index](acc, lng) }
            }) equationResult else 0
    }
}

typealias OperationPermutation = List<CalibrationOperation>

sealed class CalibrationOperation(val operation: (Long, Long) -> Long) {
    operator fun invoke(a: Long, b: Long) = operation(a, b)
    object Plus : CalibrationOperation({ a: Long, b: Long -> a + b })
    object Multiply : CalibrationOperation({ a: Long, b: Long -> a * b })
    object Concat : CalibrationOperation({ a: Long, b: Long -> "$a$b".toLong() })

    companion object {
        fun operationPermutationSequence(vararg operations: CalibrationOperation) = sequence<OperationPermutation> {
            val cache = mutableListOf<OperationPermutation>()
            val calculateCacheRange = { currentLength: Int ->
                val sectionSize = operations.size.toDouble().pow(currentLength - 1).toInt()
                val sectionStart = (1 until currentLength - 1).sumOf { operations.size.toDouble().pow(it).toInt() }
                sectionStart..(sectionStart + sectionSize - 1)
            }

            // Populate the cache with initial values for permutation length 1.
            operations.forEach { operation -> yield(listOf(operation).also { cache.add(it) }) }

            var currentLength = 2
            var offset = 0
            var cacheRange = calculateCacheRange(currentLength)
            var rotatingOperations = operations.toList()
            yieldAll(
                generateSequence {
                    if (cacheRange.count() == offset) {
                        rotatingOperations = rotatingOperations.rotated(1)
                        if (rotatingOperations.first() == operations.first()) {
                            currentLength++
                        }

                        offset = 0
                        cacheRange = calculateCacheRange(currentLength)
                    }

                    val cacheSlice = cache.slice(cacheRange)

                    return@generateSequence (cacheSlice[offset] + rotatingOperations.first()).also {
                        cache += it
                        offset++
                    } 
                }
            )
        }
    }
}

import argparse
import os


class Calibrations:
    def __init__(self, target, operators) -> None:
        self.operators = operators
        self.target = target
        self.target_found = False

    def do_math(self, numbers, operation) -> int:
        if operation == "+":
            return numbers[0] + numbers[1]
        elif operation == "*":
            return numbers[0] * numbers[1]
        elif operation == "||":
            return int(str(numbers[0]) + str(numbers[1]))

    def all_options(self, numbers, last) -> int:
        if len(numbers) < 1:
            return last
        for j in self.operators:
            # If we found our target already, abort
            # If the last value is greater than the target, abort
            if self.target_found or last > self.target:
                return
            total = self.all_options(
                numbers[1:], self.do_math((last, numbers[0]), j)
            )
            if total == self.target:
                self.target_found = True

    def process_line(self, line) -> int:
        numbers = [int(x) for x in line.split(":")[1].strip().split()]
        self.all_options(numbers[1:], numbers[0])
        if self.target_found:
            return self.target
        return 0


def main() -> None:
    path = os.path.dirname(os.path.abspath(__file__))
    parser = argparse.ArgumentParser(description="Bridge Repair")
    parser.add_argument("filename", help="Path to the input file")
    args = parser.parse_args()
    sum_of_valid = [0, 0]
    with open(f"{path}/{args.filename}", "r") as f:
        for line in f:
            line = line.strip()
            target = int(line.split(":")[0])
            for idx, ops in enumerate([["+", "*"], ["+", "*", "||"]]):
                c = Calibrations(target, ops)
                found = c.process_line(line)
                sum_of_valid[idx] += found
                if found:
                    break
    for i in range(0, 2):
        part = i + 1
        print(
            "The sum of valid calibrations for Part "
            + f"{part} is {sum(sum_of_valid[:part])}"
        )


if __name__ == "__main__":
    main()