module Puzzle where

import Data.List (intercalate)
import Data.Graph.AStar
import qualified Data.Vector as V
import qualified Data.Set as S
import Control.Monad
import Data.Maybe
import Text.Printf
import Debug.Trace

data Puzzle = Puzzle { pSize :: Int
                     , pBoard :: V.Vector Int
                     , pTwin :: Bool }
            | Start (S.Set Puzzle)
            deriving (Ord, Eq)
instance Show Puzzle where
    show p = showBoard p

showBoard :: Puzzle -> String
showBoard (Start xs) = concatMap s' $ S.toList xs
    where s' p = description p ++ showBoard p
          description p | pTwin p = "Twin:\n"
                        | otherwise = "Board:\n"
showBoard p = concatMap (showRow) (expandRows p)
    where expandRows = takeWhile (not . null) .
                       map (take size) . iterate (drop size) .
                       V.toList . pBoard
          showRow r = concatMap (printf "%4d") r ++ "\n"
          size = pSize p

column, row :: Puzzle -> Int -> Int
column p x = x `mod` pSize p
row p x = x `div` pSize p

distance :: Puzzle -> Int
distance (Start xs) = 1
distance p = V.sum . V.imap d $ pBoard p
    where d i x = if x == 0 then 0
                  else abs (c i - c x') + abs (r i - r x')
                       where x' = x - 1
          c = column p
          r = row p

findEmpty :: Puzzle -> Int
findEmpty p = fromJust $ V.findIndex (==0) (pBoard p)

swap :: Puzzle -> Int -> Int -> Maybe Puzzle
swap p e n = do
    x <- pBoard p V.!? n
    when (column p n /= column p e &&
         row p n /= row p e) Nothing
    let b' = pBoard p V.// [(e, x), (n, 0)]
    return (p { pBoard = b'})

neighbors :: Puzzle -> S.Set Puzzle
neighbors (Start xs) = xs
neighbors p = S.fromList . catMaybes $ swap' <$> ns <*> [p]
    where e = findEmpty p
          swap' n p = swap p e n
          ns = [e - 1, e + 1, e - pSize p, e + pSize p]

wrapWithTwin :: Puzzle -> Puzzle
wrapWithTwin (Start _) = error "Already wrapped"
wrapWithTwin p = Start $ S.fromList [p, p']
    where p' = p { pBoard = pBoard p V.// [(i1, e2), (i2, e1)], pTwin = True }
          (i1, i2) = (0, pSize p + 1)
          (e1, e2) = (pBoard p V.! i1, pBoard p V.! i2)

solve :: Puzzle -> [Puzzle]
solve p = case s of
            Just xs -> xs
            Nothing -> error "Could not find solution"

    where s = aStar neighbors (const (const 1)) distance ((==0) . distance) p'
          p' = wrapWithTwin p

run p = do
  putStrLn "Board:"
  putStrLn . show $ (wrapWithTwin p)
  let s = solve p
  if pTwin (s !! 0)
  then putStrLn "No solution. Showing twin:"
  else putStrLn "Solution:"
  mapM_ (putStrLn . show) (solve p)

main = do
  let example1 = Puzzle 3 (V.fromList [1, 2, 3, 4, 0, 5, 6, 7, 8]) False
  let example2 = Puzzle 3 (V.fromList [2, 1, 3, 4, 0, 5, 6, 7, 8]) False
  run example1
  run example2