content: -----------------------------------------------------------------------------
  -- |
  -- Module    : Documentation.SBV.Examples.Puzzles.Jugs
  -- Copyright : (c) Levent Erkok
  -- License   : BSD3
  -- Maintainer: erkokl@gmail.com
  -- Stability : experimental
  --
  -- Solves the classic water jug puzzle: We have 3 jugs. The capacity of the jugs are 8, 5,
  -- and 3 gallons. We begin with the 8 gallon jug full, the other two empty. We can transfer
  -- from any jug to any other, completely topping off the latter. We want to end with
  -- 4 gallons of water in the first and second jugs, and with an empty third jug. What
  -- moves should we execute in order to do so?
  -----------------------------------------------------------------------------
  
  {-# LANGUAGE DeriveAnyClass #-}
  {-# LANGUAGE DeriveGeneric  #-}
  
  {-# OPTIONS_GHC -Wall -Werror #-}
  
  module Documentation.SBV.Examples.Puzzles.Jugs where
  
  import Data.SBV
  import Data.SBV.Control
  
  import GHC.Generics(Generic)
  
  -- | A Jug has a capacity (i.e., maximum amount of water it can hold), and content, showing how much
  -- it currently has. The invariant is that content is always non-negative and is at most the capacity.
  data Jug = Jug { capacity :: Integer
                 , content  :: SInteger
                 } deriving (Generic, Mergeable)
  
  -- | Transfer from one jug to another. By definition,
  -- we transfer to fill the second jug, which may end up
  -- filling it fully, or leaving some in the first jug.
  transfer :: Jug -> Jug -> (Jug, Jug)
  transfer j1 j2 = (j1', j2')
    where empty         = literal (capacity j2) - content j2
          transferrable = empty `smin` content j1
          j1'           = j1 { content = content j1 - transferrable }
          j2'           = j2 { content = content j2 + transferrable }
  
  -- | At the beginning, we have an full 8-gallon jug, and two empty jugs, 5 and 3 gallons each.
  initJugs :: (Jug, Jug, Jug)
  initJugs = (j1, j2, j3)
    where j1 = Jug 8 8
          j2 = Jug 5 0
          j3 = Jug 3 0
  
  -- | We've solved the puzzle if 8 and 5 gallon jugs have 4 gallons each, and the third one is empty.
  solved :: (Jug, Jug, Jug) -> SBool
  solved (j1, j2, j3) = content j1 .== 4 .&& content j2 .== 4 .&& content j3 .== 0
  
  -- | Execute a bunch of moves.
  moves :: [(SInteger, SInteger)] -> (Jug, Jug, Jug)
  moves = foldl move initJugs
    where move :: (Jug, Jug, Jug) -> (SInteger, SInteger) -> (Jug, Jug, Jug)
          move (j0, j1, j2) (from, to) =
                ite ((from, to) .== (1, 2)) (let (j0', j1') = transfer j0 j1 in (j0', j1', j2))
              $ ite ((from, to) .== (2, 1)) (let (j1', j0') = transfer j1 j0 in (j0', j1', j2))
              $ ite ((from, to) .== (1, 3)) (let (j0', j2') = transfer j0 j2 in (j0', j1,  j2'))
              $ ite ((from, to) .== (3, 1)) (let (j2', j0') = transfer j2 j0 in (j0', j1,  j2'))
              $ ite ((from, to) .== (2, 3)) (let (j1', j2') = transfer j1 j2 in (j0,  j1', j2'))
              $ ite ((from, to) .== (3, 2)) (let (j2', j1') = transfer j2 j1 in (j0,  j1', j2'))
                                            (j0, j1, j2)
  
  -- | Solve the puzzle. We have:
  --
  -- >>> puzzle
  -- # of moves: 0
  -- # of moves: 1
  -- # of moves: 2
  -- # of moves: 3
  -- # of moves: 4
  -- # of moves: 5
  -- # of moves: 6
  -- # of moves: 7
  -- 1 --> 2
  -- 2 --> 3
  -- 3 --> 1
  -- 2 --> 3
  -- 1 --> 2
  -- 2 --> 3
  -- 3 --> 1
  --
  -- Here's the contents in terms of gallons after each move:
  -- (8, 0, 0)
  -- (3, 5, 0)
  -- (3, 2, 3)
  -- (6, 2, 0)
  -- (6, 0, 2)
  -- (1, 5, 2)
  -- (1, 4, 3)
  -- (4, 4, 0)
  --
  -- Note that by construction this is the minimum length solution. (Though our construction
  -- does not guarantee that it is unique.)
  puzzle :: IO ()
  puzzle = runSMT $ do
              let run i = do io $ putStrLn $ "# of moves: " ++ show (i :: Int)
                             push 1
                             ms <- mapM (const genMove) [1..i]
                             constrain $ solved $ moves ms
                             cs <- checkSat
                             case cs of
                               Unsat -> do pop 1
                                           run (i+1)
                               Sat   -> mapM_ sh ms
                               _     -> error $ "Unexpected result: " ++ show cs
              query $ run 0
    where genMove = (,) <$> freshVar_ <*> freshVar_
          sh (f, t) = do from <- getValue f
                         to   <- getValue t
                         io $ putStrLn $ show from ++ " --> " ++ show to
  