from lib import get_data, ints

data = get_data(__file__)

X, Y, Z, RADIUS = 0, 1, 2, 3


def dist(a, b):
    return sum(abs(x - y) for x, y in zip(a[:RADIUS], b[:RADIUS]))


bots = [ints(line) for line in data.splitlines()]
max_bot = max(bots, key=lambda b: b[RADIUS])
r = max_bot[3]
t = sum(1 for bot in bots if dist(bot, max_bot) <= r)
print(t)

from z3 import Int, If, Abs, Sum, IntVal, Optimize, set_param


x = Int("x")
y = Int("y")
z = Int("z")

z3zero = IntVal(0)
z3one = IntVal(1)
n_out_of_range = Int("n_out_of_range")
dist_origin = Int("dist_origin")

bots_out_of_range = [
    If(Abs(b[X] - x) + Abs(b[Y] - y) + Abs(b[Z] - z) > b[RADIUS], z3one, z3zero)
    for b in bots
]

set_param("parallel.enable", True)

opt = Optimize()
opt.add(n_out_of_range == Sum(bots_out_of_range))
opt.add(dist_origin == Abs(x) + Abs(y) + Abs(z))
opt.minimize(n_out_of_range)
opt.minimize(dist_origin)

res = opt.check()
m = opt.model()
print(m[dist_origin])