-----------------------------------------------------------------------------
-- |
-- Module    : Documentation.SBV.Examples.Puzzles.Garden
-- Copyright : (c) Levent Erkok
-- License   : BSD3
-- Maintainer: erkokl@gmail.com
-- Stability : experimental
--
-- The origin of this puzzle is Raymond Smullyan's "The Flower Garden" riddle:
--
--     In a certain flower garden, each flower was either red, yellow,
--     or blue, and all three colors were represented. A statistician
--     once visited the garden and made the observation that whatever
--     three flowers you picked, at least one of them was bound to be red.
--     A second statistician visited the garden and made the observation
--     that whatever three flowers you picked, at least one was bound to
--     be yellow.
--
--     Two logic students heard about this and got into an argument.
--     The first student said: “It therefore follows that whatever
--     three flowers you pick, at least one is bound to be blue, doesn’t
--     it?” The second student said: “Of course not!”
--
--     Which student was right, and why?
--
-- We slightly modify the puzzle. Assuming the first student is right, we use
-- SBV to show that the garden must contain exactly 3 flowers. In any other
-- case, the second student would be right.
------------------------------------------------------------------------------

{-# LANGUAGE DeriveAnyClass      #-}
{-# LANGUAGE DeriveDataTypeable  #-}
{-# LANGUAGE OverloadedStrings   #-}
{-# LANGUAGE StandaloneDeriving  #-}
{-# LANGUAGE TemplateHaskell     #-}

{-# OPTIONS_GHC -Wall -Werror #-}

module Documentation.SBV.Examples.Puzzles.Garden where

import Data.SBV
import Data.List(isSuffixOf)

-- | Colors of the flowers
data Color = Red | Yellow | Blue

-- | Make 'Color' a symbolic value.
mkSymbolicEnumeration ''Color

-- | Represent flowers by symbolic integers
type Flower = SInteger

-- | The uninterpreted function 'col' assigns a color to each flower.
col :: Flower -> SBV Color
col :: Flower -> SBV Color
col = String -> Flower -> SBV Color
forall a. Uninterpreted a => String -> a
uninterpret String
"col"

-- | Describe a valid pick of three flowers @i@, @j@, @k@, assuming
-- we have @n@ flowers to start with. Essentially the numbers should
-- be within bounds and distinct.
validPick :: SInteger -> Flower -> Flower -> Flower -> SBool
validPick :: Flower -> Flower -> Flower -> Flower -> SBool
validPick Flower
n Flower
i Flower
j Flower
k = [Flower] -> SBool
forall a. EqSymbolic a => [a] -> SBool
distinct [Flower
i, Flower
j, Flower
k] SBool -> SBool -> SBool
.&& (Flower -> SBool) -> [Flower] -> SBool
forall a. (a -> SBool) -> [a] -> SBool
sAll Flower -> SBool
ok [Flower
i, Flower
j, Flower
k]
  where ok :: Flower -> SBool
ok Flower
x = Flower -> (Flower, Flower) -> SBool
forall a. OrdSymbolic a => a -> (a, a) -> SBool
inRange Flower
x (Flower
1, Flower
n)

-- | Count the number of flowers that occur in a given set of flowers.
count :: Color -> [Flower] -> SInteger
count :: Color -> [Flower] -> Flower
count Color
c [Flower]
fs = [Flower] -> Flower
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [SBool -> Flower -> Flower -> Flower
forall a. Mergeable a => SBool -> a -> a -> a
ite (Flower -> SBV Color
col Flower
f SBV Color -> SBV Color -> SBool
forall a. EqSymbolic a => a -> a -> SBool
.== Color -> SBV Color
forall a. SymVal a => a -> SBV a
literal Color
c) Flower
1 Flower
0 | Flower
f <- [Flower]
fs]

-- | Smullyan's puzzle.
puzzle :: Goal
puzzle :: Goal
puzzle = do Flower
n <- String -> Symbolic Flower
sInteger String
"N"

            let valid :: Flower -> Flower -> Flower -> SBool
valid = Flower -> Flower -> Flower -> Flower -> SBool
validPick Flower
n

            -- Declare three existential flowers. We declare these with
            -- _modelIgnore suffix, because we don't care different assignments
            -- to them to be a different model. See 'isNonModelVar' below.
            Flower
ef1 <- String -> Symbolic Flower
forall a. SymVal a => String -> Symbolic (SBV a)
sbvExists String
"ef1_modelIgnore"
            Flower
ef2 <- String -> Symbolic Flower
forall a. SymVal a => String -> Symbolic (SBV a)
sbvExists String
"ef2_modelIgnore"
            Flower
ef3 <- String -> Symbolic Flower
forall a. SymVal a => String -> Symbolic (SBV a)
sbvExists String
"ef3_modelIgnore"

            -- Declare three universal flowers to aid in encoding the
            -- statements made by students.
            Flower
af1 <- String -> Symbolic Flower
forall a. SymVal a => String -> Symbolic (SBV a)
sbvForall String
"af1"
            Flower
af2 <- String -> Symbolic Flower
forall a. SymVal a => String -> Symbolic (SBV a)
sbvForall String
"af2"
            Flower
af3 <- String -> Symbolic Flower
forall a. SymVal a => String -> Symbolic (SBV a)
sbvForall String
"af3"

            -- Each color is represented:
            SBool -> Goal
forall (m :: * -> *). SolverContext m => SBool -> m ()
constrain (SBool -> Goal) -> SBool -> Goal
forall a b. (a -> b) -> a -> b
$ Flower -> Flower -> Flower -> SBool
valid Flower
ef1 Flower
ef2 Flower
ef3
            SBool -> Goal
forall (m :: * -> *). SolverContext m => SBool -> m ()
constrain (SBool -> Goal) -> SBool -> Goal
forall a b. (a -> b) -> a -> b
$ (Flower -> SBV Color) -> [Flower] -> [SBV Color]
forall a b. (a -> b) -> [a] -> [b]
map Flower -> SBV Color
col [Flower
ef1, Flower
ef2, Flower
ef3] [SBV Color] -> [SBV Color] -> SBool
forall a. EqSymbolic a => a -> a -> SBool
.== [SBV Color
sRed, SBV Color
sYellow, SBV Color
sBlue]

            -- Pick any three, at least one is Red
            SBool -> Goal
forall (m :: * -> *). SolverContext m => SBool -> m ()
constrain (SBool -> Goal) -> SBool -> Goal
forall a b. (a -> b) -> a -> b
$ Flower -> Flower -> Flower -> SBool
valid Flower
af1 Flower
af2 Flower
af3 SBool -> SBool -> SBool
.=> Color -> [Flower] -> Flower
count Color
Red    [Flower
af1, Flower
af2, Flower
af3] Flower -> Flower -> SBool
forall a. OrdSymbolic a => a -> a -> SBool
.>= Flower
1

            -- Pick any three, at least one is Yellow
            SBool -> Goal
forall (m :: * -> *). SolverContext m => SBool -> m ()
constrain (SBool -> Goal) -> SBool -> Goal
forall a b. (a -> b) -> a -> b
$ Flower -> Flower -> Flower -> SBool
valid Flower
af1 Flower
af2 Flower
af3 SBool -> SBool -> SBool
.=> Color -> [Flower] -> Flower
count Color
Yellow [Flower
af1, Flower
af2, Flower
af3] Flower -> Flower -> SBool
forall a. OrdSymbolic a => a -> a -> SBool
.>= Flower
1

            -- Pick any three, at least one is Blue
            SBool -> Goal
forall (m :: * -> *). SolverContext m => SBool -> m ()
constrain (SBool -> Goal) -> SBool -> Goal
forall a b. (a -> b) -> a -> b
$ Flower -> Flower -> Flower -> SBool
valid Flower
af1 Flower
af2 Flower
af3 SBool -> SBool -> SBool
.=> Color -> [Flower] -> Flower
count Color
Blue   [Flower
af1, Flower
af2, Flower
af3] Flower -> Flower -> SBool
forall a. OrdSymbolic a => a -> a -> SBool
.>= Flower
1

-- | Solve the puzzle. We have:
--
-- >>> flowerCount
-- Solution #1:
--   N = 3 :: Integer
-- This is the only solution. (Unique up to prefix existentials.)
--
-- So, a garden with 3 flowers is the only solution. (Note that we simply skip
-- over the prefix existentials and the assignments to uninterpreted function 'col'
-- for model purposes here, as they don't represent a different solution.)
flowerCount :: IO ()
flowerCount :: IO ()
flowerCount = AllSatResult -> IO ()
forall a. Show a => a -> IO ()
print (AllSatResult -> IO ()) -> IO AllSatResult -> IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< SMTConfig -> Goal -> IO AllSatResult
forall a. Provable a => SMTConfig -> a -> IO AllSatResult
allSatWith SMTConfig
z3{satTrackUFs :: Bool
satTrackUFs = Bool
False, isNonModelVar :: String -> Bool
isNonModelVar = (String
"_modelIgnore" String -> String -> Bool
forall a. Eq a => [a] -> [a] -> Bool
`isSuffixOf`)} Goal
puzzle