-----------------------------------------------------------------------------
-- |
-- Module    : Documentation.SBV.Examples.Puzzles.Tower
-- Copyright : (c) Levent Erkok
-- License   : BSD3
-- Maintainer: erkokl@gmail.com
-- Stability : experimental
--
-- Solves the tower puzzle, <http://www.chiark.greenend.org.uk/%7Esgtatham/puzzles/js/towers.html>.
-----------------------------------------------------------------------------

{-# OPTIONS_GHC -Wall -Werror #-}

module Documentation.SBV.Examples.Puzzles.Tower where

import Control.Monad
import Data.Array hiding (inRange)
import Data.SBV
import Data.SBV.Control

-------------------------------------------------------------------
-- * Modeling Towers
-------------------------------------------------------------------

-- | Count of visible towers as an array.
type Count a = Array Integer a

-- | The grid itself. The indexes are tuples, first coordinate increases as you go from
-- left to right, and the second increases as you go from top to bottom.
type Grid a = Array (Integer, Integer) a

-- | The problem has 4 counts, from top, left, bottom, and right. And the grid itself.
type Problem a = (Count a, Count a, Count a, Count a, Grid a)

-- | Example problem. Encodes:
--
-- @
--     - - 3 - - 4
--   -     2       5
--   -   2         -
--   4             -
--   2             -
--   -             2
--   3             -
--     - - 3 4 - -
-- @
problem :: Problem (Maybe Integer)
problem = (top, left, bot, right, grid)
  where build ix es = accumArray (\_ a -> a) Nothing ix [(i, Just v) | (i, v) <- es]

        top   = build  (1, 6)          [(3, 3), (6, 4)]
        left  = build  (1, 6)          [(3, 4), (4, 2), (6, 3)]
        bot   = build  (1, 6)          [(3, 3), (4, 4)]
        right = build  (1, 6)          [(1, 5), (5, 2)]
        grid  = build ((1, 1), (6, 6)) [((3, 1), 2), ((2, 2), 2)]

-- | Given a concrete partial board, turn it into a symbolic board, by filling in the
-- empty cells with symbolic variables.
symProblem :: Problem (Maybe Integer) -> Symbolic (Problem SInteger)
symProblem (t, l, b, r, g) = (,,,,) <$> fill t <*> fill l <*> fill b <*> fill r <*> fill g
 where fill :: Traversable f => f (Maybe Integer) -> Symbolic (f SInteger)
       fill = mapM (maybe free_ (pure . literal))

-------------------------------------------------------------------
-- * Counting visible towers
-------------------------------------------------------------------

-- | Given a list of tower heights, count the number of visible ones in the given order.
-- We simply keep track of the tallest we have seen so far, and increment the count for
-- each tower we see if it's taller than the tallest seen so far.
visible :: [SInteger] -> SInteger
visible = go 0 0
 where go _            visibleSofar []     = visibleSofar
       go tallestSofar visibleSofar (x:xs) = go (tallestSofar `smax` x)
                                                (ite (x .> tallestSofar) (1 + visibleSofar) visibleSofar)
                                                xs

-------------------------------------------------------------------
-- * Building constraints
-------------------------------------------------------------------

-- | Build the constraints for a given problem. We scan the elements and add the required
-- visibility counts for each row and column, viewed both in the correct order and in the backwards order.
tower :: Problem SInteger -> Symbolic ()
tower (top, left, bot, right, grid) = do
  let (minX, maxX) = bounds top
      (minY, maxY) = bounds left

  -- Constraints from top and bottom
  forM_ [minX .. maxX] $ \x -> do
      let reqT = top ! x
          reqB = bot ! x
          elts = [grid ! (x, y) | y <- [minY .. maxY]]
      mapM_ (\e -> constrain (inRange e (literal 1, literal maxY))) elts
      constrain $ distinct elts
      constrain $ reqT .== visible elts
      constrain $ reqB .== visible (reverse elts)

  -- Constraints from left and right
  forM_ [minY .. maxY] $ \y -> do
      let reqL = left  ! y
          reqR = right ! y
          elts = [grid ! (x, y) | x <- [minX .. maxX]]
      mapM_ (\e -> constrain (inRange e (literal 1, literal maxX))) elts
      constrain $ distinct elts
      constrain $ reqL .== visible elts
      constrain $ reqR .== visible (reverse elts)

-------------------------------------------------------------------
-- * Example run
-------------------------------------------------------------------

-- | Solve the puzzle descibed above. We get:
--
-- >>> example
--   1 2 3 2 2 4
-- 1 6 5 2 4 3 1 5
-- 3 3 2 5 6 1 4 2
-- 4 2 4 1 5 6 3 2
-- 2 5 3 6 1 4 2 3
-- 2 1 6 4 3 2 5 2
-- 3 4 1 3 2 5 6 1
--   3 2 3 4 2 1
example :: IO ()
example = runSMT $ do
        sp <- symProblem problem
        tower sp
        query $ do cs <- checkSat
                   case cs of
                     Unsat -> io $ putStrLn "Unsolvable"
                     Sat   -> display sp
                     _     -> error $ "Unexpected result: " ++ show cs
 where display :: Problem SInteger -> Query ()
       display (top, left, bot, right, grid) = do
          let (minX, maxX) = bounds top
              (minY, maxY) = bounds left

          -- Display top row
          io $ putStr "  "
          topVals <- forM [minX .. maxX] $ \x -> getValue (top ! x)
          io $ putStrLn $ unwords (map show topVals)

          -- Display each row, sandwiched between left/right
          forM_ [minY .. maxY] $ \y -> do
             lv <- getValue (left  ! y)
             rv <- getValue (right ! y)
             row <- forM [minX .. maxX] $ \x -> getValue (grid ! (x, y))
             io $ putStrLn $ unwords (map show (lv : row ++ [rv]))

          -- Finish with bottom row
          io $ putStr "  "
          botVals <- forM [minX .. maxX] $ \x -> getValue (bot ! x)
          io $ putStrLn $ unwords (map show botVals)