{-# LANGUAGE GeneralizedNewtypeDeriving #-}
module SliderPuzzle where

import Prelude hiding (Either(..))
import Data.Vector (Vector, (//), (!?))
import qualified Data.Vector as Vec
import Control.Arrow (second, (***))
import Control.Monad.State.Strict
import Data.Set (Set)
import qualified Data.Set as Set
import Data.PSQueue (PSQ, Binding(..))
import qualified Data.PSQueue as Queue

boardWidth, boardHeight :: Int
(boardWidth, boardHeight) = (4, 4)

data Move = Left | Right | Up | Down deriving (Eq, Ord, Show, Read)
newtype Distance = Distance Int deriving (Eq, Ord, Show, Read, Num, Integral, Enum, Real)

data Board = Board (Vector Int) [Move] Distance [Board] Distance deriving (Eq, Ord)

instance Show Board where
    show (Board vec mvs d _ dist) = "Board " ++ unwords [show vec, show mvs, show d, show dist]

getDist :: Board -> Distance
getDist (Board _ _ _ _ d) = d

idx :: (Int, Int) -> Int
idx (x, y) = y * boardWidth + x

pos :: Int -> (Int, Int)
pos i = (m, d)
    where (d, m) = i `divMod` boardWidth

getAt :: (Int, Int) -> Vector Int -> Maybe Int
getAt a@(x, y) vec | x < 0 || y < 0 || x >= boardWidth || y >= boardHeight = Nothing
                   | otherwise = vec !? idx a

swap :: (Int, Int) -> (Int, Int) -> Vector Int -> Maybe (Vector Int)
swap a b vec = do
    x <- getAt a vec
    y <- getAt b vec
    return $! vec // [(idx a, y), (idx b, x)]

moves :: Vector Int -> [(Move, Vector Int)]
moves vec = [(m, v) | (m, Just v) <- swaps]
    where (Just zero) = Vec.findIndex (== 0) vec
          (zx, zy)    = pos zero
          neighborPos = map (second $ (+ zx) *** (+ zy))
                            [(Left, (1, 0)), (Up, (0, 1)), (Right, (-1, 0)), (Down, (0, -1))]
          swaps       = map (second $ \p -> swap (zx, zy) p vec) neighborPos

manh :: (Int, Int) -> (Int, Int) -> Distance
manh (x, y) (z, w) = Distance $ abs (x - z) + abs (y - w)

distance :: Vector Int -> Distance
distance = Vec.sum . Vec.imap penalty
    where penalty i n = pos i `manh` if n == 0
                                     then (boardWidth - 1, boardHeight - 1)
                                     else pos (n - 1)

fromVecMoves :: Vector Int -> [Move] -> Distance -> Board
fromVecMoves vec mvs cnt = Board vec mvs cnt branches (distance vec)
    where branches = map (\(m, v) -> fromVecMoves v (m : mvs) (cnt + 1)) $ moves vec

fromVec :: Vector Int -> Board
fromVec initialVec = fromVecMoves initialVec [] 0

solved :: Vector Int -> Bool
solved vec = vec == Vec.fromList ([1..boardWidth * boardHeight - 1] ++ [0])

data TraverseState = TraverseState { visited :: Set (Vector Int)
                                   , queue   :: PSQ Board Distance }

solve :: Board -> Maybe [Move]
solve initial =
    evalState search (TraverseState Set.empty (Queue.singleton initial (getDist initial)))
    where search = do
              vis <- gets visited
              q   <- gets queue
              case Queue.minView q of
                  Nothing -> return Nothing
                  Just (Board vec path cnt branches _ :-> _, newQ)
                      | solved vec -> return (Just path)
                      | Set.notMember vec vis -> do
                          modify (\s -> s { queue = newQ })
                          modify (\s -> s { visited = Set.insert vec vis })
                          forM_ branches $ \b ->
                              modify $ \s ->
                                  s { queue = Queue.insert b (getDist b + cnt `div` 2) (queue s) }
                          search
                      | otherwise -> do
                          modify (\s -> s { queue = newQ })
                          search

main :: IO ()
main = print $ solve . fromVec . Vec.fromList $ [10,13,4,8,5,6,7,0,1,3,11,12,14,9,2,15]