startrek.py

from z3 import *

s = Solver()

Name, name_consts = EnumSort("Name",
                    ["Riker", "Beverly", "Data", "Troi"])
riker, beverly, data, troi = name_consts

Planet, planet_consts = EnumSort("Planet",
                    ["Betazed", "Dytallix B", "Risa", "Tellar"])
betazed, dytallix, risa, tellar = planet_consts

visited = Function("visited", Name, Planet)
s.add(Distinct([visited(name) for name in name_consts]))


Colour, colour_consts = EnumSort("Colour",
                    ["Blue", "Red", "White", "Yellow"])
blue, red, white, yellow = colour_consts
colours = Function("colours", Name, Colour)
s.add(Distinct([colours(name) for name in name_consts]))

models = Function("models", Name, IntSort())
s.add(Distinct([models(name) for name in name_consts]))
for name in name_consts:
    s.add(Or(models(name) == 5,
             models(name) == 6,
             models(name) == 7,
             models(name) == 8))

# 1. Riker has white or yellow ship models
s.add(Xor(colours(riker) == white, colours(riker) == yellow))

# 2. The crew member who visited Dytallix B and Troi have 13 ship models 
# in total. The other two crew members
# (Beverly and the one with white ship models) have 13 ship models as well.
name1_clue2, name2_clue2 = Consts("name1_clue1 name2_clue1", Name)
s.add(visited(name1_clue2) == dytallix,
      models(name1_clue2) + models(troi) == 13)
s.add(colours(name2_clue2) == white,
      models(name2_clue2) + models(beverly) == 13)

# 3. The crew member who visited Dytallix B has 1 fewer ship model than the one with the red ship models.
name1_clue3, name2_clue3 = Consts("name1_clue3 name2_clue3", Name)
s.add(visited(name1_clue3) == dytallix,
      colours(name2_clue3) == red,
      models(name1_clue3) == models(name2_clue3) - 1)

# 4. Data either visited Risa, or he has 7 ship models.
s.add(Xor(visited(data) == risa,
          models(data) == 7))

# 5. There are two more blue ship models than those owned by the crew member who visited Betazed.
name1_clue5, name2_clue5 = Consts("name1_clue5 name2_clue5", Name) 
s.add(visited(name1_clue5) == betazed,
      colours(name2_clue5) == blue,
      models(name2_clue5) == models(name1_clue5) + 2)

# 6. The crew member who visited Tellar doesn't own white ship models.
name_clue6 = Const("name_clue6", Name)
s.add(visited(name_clue6) == tellar,
      colours(name_clue6) != white)

# 7. Only the crew member with 7 ship models has the same initial as with the place they visited.
name_clue7 = Const("name_clue7", Name)
s.add(AtMost(visited(data) == dytallix, 
          visited(troi) == tellar,
          visited(riker) == risa,
          visited(beverly) == betazed,
      1))
s.add(Implies(name_clue7 == data, visited(data) == dytallix),
      Implies(name_clue7 == troi, visited(troi) == tellar),
      Implies(name_clue7 == riker, visited(riker) == risa),
      Implies(name_clue7 == beverly, visited(beverly) == betazed))
        
print(s.check())      
if s.check() == sat:
    m = s.model()
    for name in name_consts:
        print("{}: {} colour, {} models, {} visited"
          .format(name,
              m.eval(colours(name)),
              m.eval(models(name)),
              m.eval(visited(name))))