import com.svetylkovo.rojo.Rojo
import com.svetylkovo.rojo.annotations.Group
import com.svetylkovo.rojo.annotations.Regex
import java.io.File

typealias InstructionFunc = (Int, Int, Int) -> Unit

class WristDevice {
    val registers = longArrayOf(0, 0, 0, 0)

    fun setRegisters(value: LongArray) {
        value.copyInto(registers)
    }

    val instructions: List<InstructionFunc> = listOf(
        this::addr,
        this::addi,
        this::mulr,
        this::muli,
        this::banr,
        this::bani,
        this::borr,
        this::bori,
        this::setr,
        this::seti,
        this::gtir,
        this::gtri,
        this::gtrr,
        this::eqir,
        this::eqri,
        this::eqrr
    )

    val opToInstruction = mutableMapOf<Int, InstructionFunc>()

    /*
     * Instructions
     */

    fun addr(regA: Int, regB: Int, output: Int) {
        registers[output] = registers[regA] + registers[regB]
    }

    fun addi(regA: Int, valueB: Int, output: Int) {
        registers[output] = registers[regA] + valueB
    }

    fun mulr(regA: Int, regB: Int, output: Int) {
        registers[output] = registers[regA] * registers[regB]
    }

    fun muli(regA: Int, valueB: Int, output: Int) {
        registers[output] = registers[regA] * valueB
    }

    fun banr(regA: Int, regB: Int, output: Int) {
        registers[output] = registers[regA] and registers[regB]
    }

    fun bani(regA: Int, valueB: Int, output: Int) {
        registers[output] = registers[regA] and valueB.toLong()
    }

    fun borr(regA: Int, regB: Int, output: Int) {
        registers[output] = registers[regA] or registers[regB]
    }

    fun bori(regA: Int, valueB: Int, output: Int) {
        registers[output] = registers[regA] or valueB.toLong()
    }

    fun setr(regA: Int, ignored: Int, output: Int) {
        registers[output] = registers[regA]
    }

    fun seti(valueA: Int, ignored: Int, output: Int) {
        registers[output] = valueA.toLong()
    }

    fun gtir(valueA: Int, regB: Int, output: Int) {
        registers[output] = if (valueA > registers[regB]) 1 else 0
    }

    fun gtri(regA: Int, valueB: Int, output: Int) {
        registers[output] = if (registers[regA] > valueB) 1 else 0
    }

    fun gtrr(regA: Int, regB: Int, output: Int) {
        registers[output] = if (registers[regA] > registers[regB]) 1 else 0
    }

    fun eqir(valueA: Int, regB: Int, output: Int) {
        registers[output] = if (valueA.toLong() == registers[regB]) 1 else 0
    }

    fun eqri(regA: Int, valueB: Int, output: Int) {
        registers[output] = if (registers[regA] == valueB.toLong()) 1 else 0
    }

    fun eqrr(regA: Int, regB: Int, output: Int) {
        registers[output] = if (registers[regA] == registers[regB]) 1 else 0
    }

    /*
     * End of Instructions
     */

    fun getMatchingInstructions(sample: Sample) = instructions.filter { instruction ->
        setRegisters(sample.input.toLongArray())
        instruction(
            sample.instruction[1].toInt(),
            sample.instruction[2].toInt(),
            sample.instruction[3].toInt()
        )
        registers.contentEquals(sample.output.toLongArray())
    }

    fun figureOutOpNumbers(samples: List<Sample>) {
        val matchingOpsMap = samples.associateBy { it.instruction[0].toInt() }
            .mapValues { getMatchingInstructions(it.value) }

        while (opToInstruction.size != 16) {
            matchingOpsMap.forEach { op, functions ->
                val remaining = functions - opToInstruction.values
                if (remaining.size == 1) {
                    opToInstruction += op to remaining.single()
                }
            }
        }
    }

    fun execute(instruction: Instruction) {
        opToInstruction[instruction.op.toInt()]?.let { instructionFunction ->
            instructionFunction(
                instruction.inputA.toInt(),
                instruction.inputB.toInt(),
                instruction.outputC.toInt()
            )
        }
    }

}

@Regex("Before:(.+)\\n(.+)\\nAfter: (.+)")
class Sample {
    @Group(1)
    @Regex("\\d+")
    var input: List<Long> = emptyList()
    @Group(2)
    @Regex("\\d+")
    var instruction: List<Long> = emptyList()
    @Group(3)
    @Regex("\\d+")
    var output: List<Long> = emptyList()
}

@Regex("(\\d+).+(\\d+).+(\\d+).+(\\d+)")
class Instruction {
    @Group(1)
    var op: Long = 0
    @Group(2)
    var inputA: Long = 0
    @Group(3)
    var inputB: Long = 0
    @Group(4)
    var outputC: Long = 0
}

fun main() {

    val inputIterator = File("input16.txt").readLines().iterator()

    //A
    val inputA = inputIterator.asSequence().take(3109).joinToString("\n")

    val samples = Rojo.of(Sample::class.java).matchList(inputA)

    val wristDeviceA = WristDevice()

    val resultA = samples.map { wristDeviceA.getMatchingInstructions(it).count() }
        .count { it >= 3 }

    println(resultA)

    //B
    val inputB = inputIterator.asSequence().joinToString("\n")
    val instructions = Rojo.of(Instruction::class.java).matchList(inputB)

    val wristDeviceB = WristDevice()
    wristDeviceB.figureOutOpNumbers(samples)

    wristDeviceB.setRegisters(longArrayOf(0, 0, 0, 0))
    instructions.forEach { wristDeviceB.execute(it) }

    println(wristDeviceB.registers.joinToString(","))
}