; Classic zebra logic puzzle solution
; Author: Bruno Wolff III
; Last updated: 2020-08-05

; You may use/modify this code for any purpose

(set-option :auto-config false)
(set-option :produce-proofs true)
(set-option :produce-models true)

; Declare order to set size for arrays and to handle ordering relations of houses
(declare-datatypes () ((order one two three four five)))
(declare-const ordering (Array order Int))

(declare-datatypes () ((nationallity Englishman Spaniard Ukrainian Norwegian Japanese)))
(declare-const nationallities (Array order nationallity))
(assert (forall ((o1 order) (o2 order)) (=> (distinct o1 o2) (distinct (select nationallities o1) (select nationallities o2)))))

; Possible smokes
(declare-datatypes () ((smoke Old_Gold Kools Chesterfields Lucky_Strike Parliaments)))
(declare-const smokes (Array order smoke))
(assert (forall ((o1 order) (o2 order)) (=> (distinct o1 o2) (distinct (select smokes o1) (select smokes o2)))))

; Possible houses
(declare-datatypes () ((color red green ivory yellow blue)))
(declare-const houses (Array order color))
(assert (forall ((o1 order) (o2 order)) (=> (distinct o1 o2) (distinct (select houses o1) (select houses o2)))))

; possible pets
(declare-datatypes () ((pet dog snails fox horse zebra)))
(declare-const pets (Array order pet))
(assert (forall ((o1 order) (o2 order)) (=> (distinct o1 o2) (distinct (select pets o1) (select pets o2)))))

; possible drinks
(declare-datatypes () ((drink tea coffee milk orange_juice water)))
(declare-const drinks (Array order drink))
(assert (forall ((o1 order) (o2 order)) (=> (distinct o1 o2) (distinct (select drinks o1) (select drinks o2)))))

; There are five houses.
(assert (and (= (select ordering one) 1) (= (select ordering two) 2) (= (select ordering three) 3) (= (select ordering four) 4) (= (select ordering five) 5)))

; The Englishman lives in the red house.
(assert (forall ((o order)) (=> (= (select nationallities o) Englishman) (= (select houses o) red))))

; The Spaniard owns the dog.
(assert (forall ((o order)) (=> (= (select nationallities o) Spaniard) (= (select pets o) dog))))

; Coffee is drunk in the green house.
(assert (forall ((o order)) (=> (= (select drinks o) coffee) (= (select houses o) green))))

; The Ukrainian drinks tea.
(assert (forall ((o order)) (=> (= (select nationallities o) Ukrainian) (= (select drinks o) tea))))

; The green house is immediately to the right of the ivory house.
(assert (forall ((o1 order) (o2 order)) (=> (and (= (select houses o1) green) (= (select houses o2) ivory)) (= (select ordering o1) (+ (select ordering o2) 1)))))

; The Old Gold smoker owns snails.
(assert (forall ((o order)) (=> (= (select smokes o) Old_Gold) (= (select pets o) snails))))

; Kools are smoked in the yellow house.
(assert (forall ((o order)) (=> (= (select smokes o) Kools) (= (select houses o) yellow))))

; Milk is drunk in the middle house.
(assert (= (select drinks three) milk))

; The Norwegian lives in the first house.
(assert (= (select nationallities one) Norwegian))

; The man who smokes Chesterfields lives in the house next to the man with the fox.
(assert (forall ((o1 order) (o2 order)) (=> (and (= (select smokes o1) Chesterfields) (= (select pets o2) fox)) (or (= (select ordering o1) (+ (select ordering o2) 1)) (= (select ordering o1) (- (select ordering o2) 1))))))

; Kools are smoked in the house next to the house where the horse is kept.
(assert (forall ((o1 order) (o2 order)) (=> (and (= (select smokes o1) Kools) (= (select pets o2) horse)) (or (= (select ordering o1) (+ (select ordering o2) 1)) (= (select ordering o1) (- (select ordering o2) 1))))))

; The Lucky Strike smoker drinks orange juice.
(assert (forall ((o order)) (=> (= (select smokes o) Lucky_Strike) (= (select drinks o) orange_juice))))

; The Japanese smokes Parliaments.
(assert (forall ((o order)) (=> (= (select nationallities o) Japanese) (= (select smokes o) Parliaments))))

; The Norwegian lives next to the blue house.
(assert (forall ((o1 order) (o2 order)) (=> (and (= (select nationallities o1) Norwegian) (= (select houses o2) blue)) (or (= (select ordering o1) (+ (select ordering o2) 1)) (= (select ordering o1) (- (select ordering o2) 1))))))

(echo "")
(echo "Solution 1")
(echo "")
(check-sat)
(echo "")
(eval one)
(eval (select ordering one))
(eval (select nationallities one))
(eval (select houses one))
(eval (select smokes one))
(eval (select pets one))
(eval (select drinks one))
(echo "")
(eval two)
(eval (select ordering two))
(eval (select nationallities two))
(eval (select houses two))
(eval (select smokes two))
(eval (select pets two))
(eval (select drinks two))
(echo "")
(eval three)
(eval (select ordering three))
(eval (select nationallities three))
(eval (select houses three))
(eval (select smokes three))
(eval (select pets three))
(eval (select drinks three))
(echo "")
(eval four)
(eval (select ordering four))
(eval (select nationallities four))
(eval (select houses four))
(eval (select smokes four))
(eval (select pets four))
(eval (select drinks four))
(echo "")
(eval five)
(eval (select ordering five))
(eval (select nationallities five))
(eval (select houses five))
(eval (select smokes five))
(eval (select pets five))
(eval (select drinks five))
(echo "")

; Unfortunately z3 doesn't provide a way to extract a model and use it while reading smt-lib formatted input. So to check for other solutions requires and constraints and running again until there aren't any more.

(assert-not (and

  (= (select nationallities one) Norwegian)
  (= (select houses one) yellow)
  (= (select smokes one) Kools)
  (= (select pets one) fox)
  (= (select drinks one) water)

  (= (select nationallities two) Ukrainian)
  (= (select houses two) blue)
  (= (select smokes two) Chesterfields)
  (= (select pets two) horse)
  (= (select drinks two) tea)

  (= (select nationallities three) Englishman)
  (= (select houses three) red)
  (= (select smokes three) Old_Gold)
  (= (select pets three) snails)
  (= (select drinks three) milk)

  (= (select nationallities four) Spaniard)
  (= (select houses four) ivory)
  (= (select smokes four) Lucky_Strike)
  (= (select pets four) dog)
  (= (select drinks four) orange_juice)

  (= (select nationallities five) Japanese)
  (= (select houses five) green)
  (= (select smokes five) Parliaments)
  (= (select pets five) zebra)
  (= (select drinks five) coffee)

))

(echo "Solution 2")
(echo "")
(check-sat)
(echo "")

(exit)