ZebraPuzzle.hs

module ZebraPuzzle (Resident(..), Solution(..), solve) where

import Prelude hiding (product, not, and, (&&))

import Data.List (find)
import Data.Maybe (fromJust)
import Data.Ix (Ix, range)
import qualified Data.Array as A
import Control.Monad.Trans.State.Lazy (StateT)
import System.IO.Unsafe (unsafePerformIO)

import Ersatz hiding (Solution)
import Ersatz.Relation


data Solution = Solution { waterDrinker :: Resident
                         , zebraOwner :: Resident
                         } deriving (Eq, Show)


data House = First | Second | Third | Fourth | Fifth
  deriving (Eq, Ord, Show, Enum, Bounded, Ix)

data Resident = Englishman | Spaniard | Ukrainian | Norwegian | Japanese
  deriving (Eq, Ord, Show, Enum, Bounded, Ix)

data Color = Red | Green | Ivory | Yellow | Blue
  deriving (Eq, Ord, Show, Enum, Bounded, Ix)

data Drink = Coffee | Tea | Milk | OrangeJuice | Water
  deriving (Eq, Ord, Show, Enum, Bounded, Ix)

data Pet = Dog | Snails | Fox | Horse | Zebra
  deriving (Eq, Ord, Show, Enum, Bounded, Ix)

data Cigarettes = OldGold | Kool | Chesterfield | LuckyStrike | Parliament
  deriving (Eq, Ord, Show, Enum, Bounded, Ix)



-- UTILITIES ------------------------------------------------------------------

fullBounds :: Bounded a => (a, a)
fullBounds = (minBound, maxBound)


succRelation :: (Ix a, Enum a) => ((a, a), (a, a)) -> Relation a a
succRelation bnd =
  build bnd $ flip map (range bnd) $ \(i,j) ->
    -- we need the `i < j` check since `succ` fails on maxBound
    ((i,j), if i < j && succ i == j then true else false)

rightOf, leftOf, neighborOf :: Relation House House
leftOf = succRelation fullBounds
rightOf = mirror leftOf
neighborOf = union leftOf rightOf

bijection :: (Ix a, Ix b) => Relation a b -> Bit
bijection rel = regular_in_degree 1 rel && regular_out_degree 1 rel

-- just an infix alias, similar to composition in Control.Category
(>>>) :: (Ix a, Ix b, Ix c) => Relation a b -> Relation b c -> Relation a c
(>>>) = product


-- PROBLEM DESCRIPTION ---------------------------------------------------------

type Solve = StateT SAT IO

problem :: Solve (Relation Resident Drink, Relation Resident Pet)
problem = do
  resident <- relation fullBounds :: Solve (Relation House Resident)
  color    <- relation fullBounds :: Solve (Relation House Color)
  drink    <- relation fullBounds :: Solve (Relation House Drink)
  pet      <- relation fullBounds :: Solve (Relation House Pet)
  cigs     <- relation fullBounds :: Solve (Relation House Cigarettes)

  -- Each of the five houses is painted a different color, and their
  -- inhabitants are of different national extractions, own different pets,
  -- drink different beverages and smoke different brands of cigarettes.
  assert . and $
    [ bijection resident
    , bijection color
    , bijection pet
    , bijection drink
    , bijection cigs
    ]

  -- 2. The Englishman lives in the red house.
  assert $ mirror resident >>> color ! (Englishman, Red)
  -- 3. The Spaniard owns the dog.
  assert $ mirror resident >>> pet ! (Spaniard, Dog)
  -- 4. Coffee is drunk in the green house.
  assert $ mirror drink >>> color ! (Coffee, Green)
  -- 5. The Ukrainian drinks tea.
  assert $ mirror resident >>> drink ! (Ukrainian, Tea)
  -- 6. The green house is immediately to the right of the ivory house.
  assert $ mirror color >>> rightOf >>> color ! (Green, Ivory)
  -- 7. The Old Gold smoker owns snails.
  assert $ mirror cigs >>> pet ! (OldGold, Snails)
  -- 8. Kools are smoked in the yellow house.
  assert $ mirror cigs >>> color ! (Kool, Yellow)
  -- 9. Milk is drunk in the middle house.
  assert $ drink ! (Third, Milk)
  -- 10. The Norwegian lives in the first house.
  assert $ resident ! (First, Norwegian)
  -- 11. The man who smokes Chesterfields lives in the house next to the man with the fox.
  assert $ mirror cigs >>> neighborOf >>> pet ! (Chesterfield, Fox)
  -- 12. Kools are smoked in the house next to the house where the horse is kept
  assert $ mirror cigs >>> neighborOf >>> pet ! (Kool, Horse)
  -- 13. The Lucky Strike smoker drinks orange juice.
  assert $ mirror cigs >>> drink ! (LuckyStrike, OrangeJuice)
  -- 14. The Japanese smokes Parliaments.
  assert $ mirror resident >>> cigs ! (Japanese, Parliament)
  -- 15. The Norwegian lives next to the blue house.
  assert $ mirror resident >>> neighborOf >>> color ! (Norwegian, Blue)

  return (mirror resident >>> drink, mirror resident >>> pet)


-- RUNNING THE SOLVER ---------------------------------------------------------

solveIO :: IO (Maybe Solution)
solveIO = do
  (_, mRes) <- solveWith minisat problem  -- requires `minisat` in PATH
  return $ do
    (drink, pet) <- mRes
    -- Which of the residents drinks water?
    -- Who owns the zebra?
    Solution <$> residentWith drink Water <*> residentWith pet Zebra
  where
    residentWith :: Ix a => A.Array (Resident, a) Bool -> a -> Maybe Resident
    residentWith rel a = find (\r -> rel A.! (r, a)) (range fullBounds)