from z3 import And, Array, If, Int, IntSort, Or, Solver, Store

codes = open("inputs/2.txt").read().strip().split(",")
codes = list(map(int, codes))
N = len(codes)

s = Solver()

# Initialize array replacing position 1 and 2 with variables a and b
A = Array("A", IntSort(), IntSort())
a = Int("a")
b = Int("b")
for i, c in enumerate(codes):
    if i == 1:
        A = Store(A, i, a)
    elif i == 2:
        A = Store(A, i, b)
    else:
        A = Store(A, i, codes[i])
# As a hack we also store the instruction pointer at position N
A = Store(A, N, 0)

# Each iteration calculates A at timestep t
# Since A is both our instructions and instruction pointer, we can update both simultaneously
for t in range(0, N, 4):
    ip = A[N]
    # Check if it's a valid op
    op = A[ip]
    s.add(Or(op == 1, op == 2, op == 99))
    # If not stopped, the indices also have to be valid
    stopped = op == 99
    validIndices = And(
        A[ip + 1] >= 0,
        A[ip + 1] < N,
        A[ip + 2] >= 0,
        A[ip + 2] < N,
        A[ip + 3] >= 0,
        A[ip + 3] < N,
    )
    s.add(Or(stopped, validIndices))
    # Update A with the new A that has A[N] += 4 and the result of the op
    # If op is not 1 or 2, the new A is just the old A (where instruction pointer isn't incremented either)
    A = If(
        op == 1,
        Store(Store(A, N, ip + 4), A[ip + 3], A[A[ip + 1]] + A[A[ip + 2]]),
        If(
            op == 2,
            Store(Store(A, N, ip + 4), A[ip + 3], A[A[ip + 1]] * A[A[ip + 2]]),
            # else no-op
            A,
        ),
    )

s.add(A[A[N]] == 99)  # Program is stopped
s.add(A[0] == 19690720)  # The desired target value

print("isSat", s.check())
m = s.model()
print("model", m)
a_long = m[a].as_long()
b_long = m[b].as_long()
print("a", a_long)
print("b", b_long)
print("Part 2", 100 * a_long + b_long)