# There are five house addresses: 1, 2, 3, 4, 5
# There are five distinct people: Axl, Bev, Che, Dez, Eva
# There are five distinct house colors: Red, Blue, Green, Yellow, and Purple.
# There are five distinct beings: Earthling, Martian, Jovian (from Jupiter), Saturnian, Venusian.
# There are five distinct beverages: Water, Tea, Milk, Juice, Coffee.
# There are five distinct pets: Fox, Horse, Snail, Dog, Zebra.


# Each person lives in one house, has one name, is from one planet, drinks one beverage, and has one pet.
# Each pet belongs to exactly one person. 

# The Earthling lives in the Red house.
# The Martian has a pet Dog.
# Coffee is drunk in the Green house.
# The Jovian drinks Tea.
# The Green house is immediately to the right of the Purple house.
# Axl has a pet Snail.
# Bev lives in the Yellow house.
# Milk is drunk in the middle house.
# The Saturnian lives in the first house.
# Che lives in the house next to the person with the Fox.
# Bev lives in the house next to the house where the Horse is kept.
# Eva drinks Juice.
# Dez is from Venus.
# The Saturnian lives next to the Blue house.

from z3 import *

# names
Axl = Int('Axl')
Bev = Int('Bev')
Che = Int('Che')
Dez = Int('Dez')
Eva = Int('Eva')
names = [Axl, Bev, Che, Dez, Eva]

# house colors
Red = Int('Red')
Blue = Int('Blue')
Green = Int('Green')
Yellow  = Int('Yellow ')
Purple = Int('Purple')
colors = [Red, Blue, Green, Yellow, Purple]

# planets
Earth = Int('Earth')
Mars = Int('Mars')
Jupiter = Int('Jupiter')
Saturn = Int('Saturn')
Venus = Int('Venus')
planets = [Earth, Mars, Jupiter, Saturn, Venus]

# Beverages
Water = Int('Water')
Tea = Int('Tea')
Milk = Int('Milk')
Juice = Int('Juice')
Coffee = Int('Coffee')
beverages = [Water, Tea, Milk, Juice, Coffee]

# Pets
Fox = Int('Fox')
Horse = Int('Horse')
Snail = Int('Snail')
Dog = Int('Dog')
Zebra = Int('Zebra')
pets = [Fox, Horse, Snail, Dog, Zebra]

var_lists = [names, colors, planets, beverages, pets]

different_names = Distinct(Axl, Bev, Che, Dez, Eva)
different_colors = Distinct(Red, Blue, Green, Yellow, Purple)
different_planets = Distinct(Earth, Mars, Jupiter, Saturn, Venus)
different_beverages = Distinct(Water, Tea, Milk, Juice, Coffee)
different_pets = Distinct(Fox, Horse, Snail, Dog, Zebra)

ranges = [And(1 <= x, x <= 5)  for var_list in var_lists for x in var_list]


constraints = []
# The Earthling lives in the Red house.
constraints.append(Earth == Red)

# The Martian has a pet Dog.
constraints.append(Mars == Dog)

# Coffee is drunk in the Green house.
constraints.append(Coffee == Green)

# The Jovian drinks Tea.
constraints.append(Jupiter == Tea)

# The Green house is immediately to the right of the Purple house.
constraints.append(Green == Purple + 1)

# Axl has a pet Snail.
constraints.append(Axl == Snail)

# Bev lives in the Yellow house.
constraints.append(Bev == Yellow)

# Milk is drunk in the middle house.
constraints.append(Milk == 3)

# The Saturnian lives in the first house.
constraints.append(Saturn == 1)

# Che lives in the house next to the person with the Fox.
constraints.append(Or(Che == Fox + 1, Che == Fox - 1))

# Bev lives in the house next to the house where the Horse is kept.
constraints.append(Or(Bev == Horse + 1, Bev == Horse - 1))

# Eva drinks Juice.
constraints.append(Eva == Juice)

# Dez is from Venus.
constraints.append(Dez == Venus)

# The Saturnian lives next to the Blue house.
constraints.append(Or(Saturn == Blue + 1, Saturn == Blue - 1))

s = Solver();
s.add(different_beverages)
s.add(different_colors)
s.add(different_names)
s.add(different_pets)
s.add(different_planets)

s.add(ranges)
s.add(constraints)

# how to negate a model and check uniqueness? 

s.check()
m = s.model()

# collect the variables from the model m
model_vars = [v for v in m]
# collect the values of all of the variables
model_vals = [m[v] for v in model_vars]

# take disjunction of negations of variable assignments 
blocked_model = Or([(var() != val) for (var,val) in zip(model_vars, model_vals)])

# add the blocked clause
s.add(blocked_model)

# now check sat again
s.check()