import Foundation

/*
0 adv: A / 2^OP -> A
1 bxl: B xor LIT -> B
2 bst: OP % 8 -> B
3 jnz: jump LIT if A != 0
4 bxc: B xor C -> B (ignore op/lit)
5 out: OP % 8 -> screen
6 bdv: A / 2^OP -> B
7 cdv: A / 2^OP -> C

2,4, 1,1, 7,5, 1,5, 4,0, 5,5, 0,3, 3,0
bst 4 // B = A % 8
bxl 1 // B = B ^ 1
cdv 5 // C = A >> B
bxl 5 // B = B ^ 5
bxc 0 // B = B ^ C
out 5 // print(B % 8)
adv 3 // A = A >> 3
jnz 0 // loop while A != 0
*/

extension RangeReplaceableCollection where Self: StringProtocol {
    func padLeft(to length: Int, using e: Element = " ") -> SubSequence {
        return repeatElement(e,
            count: Swift.max(0, length-count)) +
                suffix(Swift.max(count, count-length)
        )
    }
}

/*
printout 2
B1 = A%8 | B2  | C3   | B4  | B5 = B4^C3 = xxxxx010 | A?
---------+-----+------+-----+-----------------------+------------
     000 | 001 | A>>1 | 100 | 100^(A>>1) = xxxxx010 | ...xxxx001x (❌)
     001 | 000 | A    | 101 | 101^A      = xxxxx010 | ...xxxxx111 (❌)
     010 | 011 | A>>3 | 110 | 110^(A>>3) = xxxxx010 | ...xx100010
     011 | 010 | A>>2 | 111 | 111^(A>>2) = xxxxx010 | ...xxx101xx (❌)
     100 | 101 | A>>5 | 000 | 000^(A>>5) = xxxxx010 | ...010xx100
     101 | 100 | A>>4 | 001 | 001^(A>>4) = xxxxx010 | ...x011x101
     110 | 111 | A>>7 | 010 | 010^(A>>7) = xxxxx010 | .000xxxx110
     111 | 110 | A>>6 | 011 | 011^(A>>6) = xxxxx010 | ..001xxx111
*/

struct Register : Comparable, CustomStringConvertible {
    let value: Int
    let mask: Int
    var description: String {
        var s: String = ""
        var (v, m) = (value, mask)
        while v != 0 || m != 0 {
            s = ((m&1 == 0) ? "?" : String(v&1)) + s
            v >>= 1
            m >>= 1
        }
        return String(s.padLeft(to: 16, using: "_"))
    }
    static func < (lhs: Register, rhs: Register) -> Bool {
        return (lhs.value & lhs.mask) < (rhs.value & rhs.mask)
    }
    static func == (lhs: Register, rhs: Register) -> Bool {
        return (lhs.value & lhs.mask) == (rhs.value & rhs.mask)
    }
    var dec: Int { return value & mask }

    func combine(with b: Register) -> Register? {
        if self.value & self.mask & b.mask == b.value & self.mask & b.mask {
            return Register(
                value: (self.value & self.mask) | (b.value & b.mask),
                mask: self.mask | b.mask
            )
        }
        return nil
    }
    static func >> (lhs: Register, rhs: Int) -> Register {
        return Register(value: lhs.value >> rhs, mask: lhs.mask >> rhs)
    }
    static func << (lhs: Register, rhs: Int) -> Register {
        return Register(value: lhs.value << rhs, mask: lhs.mask << rhs)
    }
}

func possibleAs(fromB5 b5Suffix: Int) -> [Register] {
    let choices: [Register] = (0b000...0b111).compactMap { aSuffix in
        let aShift = aSuffix ^ 1  // Replace this with first XOR literal
        let b4 = aShift ^ 5       // Replace this with second XOR literal
        let a = Register(value: (b4^b5Suffix) << aShift, mask: 0b111 << aShift)
        return a.combine(with: Register(value: aSuffix, mask: 0b111))
    }
    return choices.sorted()
}

func findA(from a: Register, _ outputs: [Int]) -> [Register] {
    if outputs.count == 0 { return [a] }
    let As = possibleAs(fromB5: outputs[0]).compactMap { a.combine(with: $0) }
    return As.flatMap { a in 
        findA(from: a >> 3, Array(outputs.dropFirst())).compactMap { nextA in
            a.combine(with: nextA << 3)
        }
    }.sorted()
}

let prog = [2,4, 1,1, 7,5, 1,5, 4,0, 5,5, 0,3, 3,0]
let res = findA(from: Register(value: 0, mask: 0), prog)
res.forEach { print($0) }
print(res[0].dec)