bucket_puzzle.hs

-- https://gist.github.com/sordina/ea8c66ddcc467f77b58b0af0b5628435

module Lib ( main ) where

import Data.Tree
import Data.List
import System.Environment

type Capacity = Int
type Level    = Int
type Index    = Int
data State    = State { buckets :: [Bucket], actions :: [Action] } deriving (Show, Eq)
data Bucket   = Bucket { capacity :: Capacity, level :: Level }    deriving (Show, Eq)
data Action   = Transfer Index Index | Empty Index | Fill Index    deriving (Show, Eq)

options :: State -> [State]
options (State bs as) = nextStates
  where
  possible   = possibilities bs
  outcomes   = map (applyTo bs) possible
  nextStates = zipWith (makeState as) outcomes possible

applyTo :: [Bucket] -> Action -> [Bucket]
applyTo bs (Empty    b1   ) = replace b1 empty bs
applyTo bs (Fill     b1   ) = replace b1 fill  bs
applyTo bs (Transfer x1 x2) = transfer bs x1 x2

-- Most complicated function...
transfer :: [Bucket] -> Index -> Index -> [Bucket]
transfer bs x1 x2 | l1 <= gap = replace x2 (add l1) $ replace x1 empty bs
                  | l1 >  gap = replace x2 fill     $ replace x1 (remove gap) bs
                  | otherwise = bs
  where
  b1  = bs !! x1
  b2  = bs !! x2
  l1  = level b1
  l2  = level b2
  c2  = capacity b2
  gap = c2 - l2

-- prop_foo :: Bool
-- prop_foo = applyTo [Bucket 3 3, Bucket 5 0] (Transfer 0 1) == [Bucket 3 0, Bucket 5 3]

add :: Int -> Bucket -> Bucket
add v (Bucket c l) = Bucket c (l + v)

remove :: Int -> Bucket -> Bucket
remove v (Bucket c l) = Bucket c (l - v)

replace :: Int -> (a -> a) -> [a] -> [a]
replace n f = zipWith at [0..]
  where
  at x a | n == x = f a
         | otherwise = a

empty :: Bucket -> Bucket
empty (Bucket c _l) = Bucket c 0

fill :: Bucket -> Bucket
fill (Bucket c _l) = Bucket c c

makeState :: [Action] -> [Bucket] -> Action -> State
makeState as bs a = State bs (a : as)

possibilities :: [Bucket] -> [Action]
possibilities bs = map Empty empties ++ map Fill fills ++ transfers
  where
  elem'     = flip elem
  empties   = indexes (not . isEmpty) bs
  fills     = indexes (not . isFull) bs
  pairs     = [(pred f, pred t) | f <- [1.. length bs], t <- [1.. length bs]]
  transfers = map (uncurry Transfer)
            $ filter (elem' empties . fst)
            $ filter (elem' fills . snd)
            $ filter (uncurry (/=)) pairs

indexes :: (a -> Bool) -> [a] -> [Int]
indexes f = map fst . filter snd . zip [0..] . map f

isEmpty :: Bucket -> Bool
isEmpty (Bucket _c l) = l <= 0

isFull :: Bucket -> Bool
isFull (Bucket c l) = l >= c

buildTree :: (State -> [State]) -> State -> Tree State
buildTree f s = Node s (map (buildTree f) (f s))

powerset :: [a] -> [[a]]
powerset [] = [[]]
powerset (x:xs) = [x:ps | ps <- powerset xs] ++ powerset xs

powerset' :: [a] -> [[a]]
powerset' = filter (not . null) . powerset

-- Prevent cycles
prune :: Level -> Tree State -> Tree State
prune target (Node x fs)
  | solved target x = Node x []
  | otherwise = Node x (map (prune target) fs)

initialState :: [Level] -> State
initialState bs = State (map (`Bucket` 0) bs) []

solved :: Level -> State -> Bool
solved target (State bs _as) = any ((== target) . sum . map level) (powerset' bs)

help :: IO ()
help = putStrLn "Usage: bucket-solver TARGET CAPACITY*"

main :: IO ()
main = do
  args <- getArgs
  case args of
    (t : bs@(_:_)) -> mapM_ print $ take 10 $ nub $ filter (solved (read t)) $ concat $ levels $ prune (read t) $ buildTree options (initialState (map read bs))
    _              -> help