{-# LANGUAGE ParallelListComp, RecordWildCards #-}
module Main where

import Data.Char
import Data.List (sort)
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Set (Set)
import qualified Data.Set as Set
import System.Environment (getArgs)

type Slots a = Map a (Set Int)
data Puzzle = Puzzle { puzzleGrid :: Map Pos Int,
                       puzzleOpen :: Set Pos,
                       puzzleCols :: Slots Int,
                       puzzleRows :: Slots Int,
                       puzzleBlocks :: Slots Pos }
type Pos = (Int, Int)

posBlock (x, y) = (x `div` 3, y `div` 3)

puzzleInsert pos@(x, y) n (Puzzle {..}) =
    Puzzle (Map.insert pos n puzzleGrid) (pos `Set.delete` puzzleOpen)
           (removeSlot x puzzleCols) (removeSlot y puzzleRows)
           (removeSlot (posBlock pos) puzzleBlocks)
  where
    removeSlot :: Ord a => a -> Slots a -> Slots a
    removeSlot = Map.adjust (Set.delete n)

find :: Ord a => a -> Slots a -> Set Int
find = Map.findWithDefault undefined

availableNumbers pos@(x, y) (Puzzle {..}) =
    find x puzzleCols `Set.intersection` find y puzzleRows
    `Set.intersection` find (posBlock pos) puzzleBlocks

forwardCheck pos@(x, y) p = check (gridCols x) && check (gridRows y)
                            && check (gridBlocks (posBlock pos))
  where
    check = not . any (\pos -> if pos `Map.member` puzzleGrid p then False
                               else Set.null (availableNumbers pos p))

numValuesRuledOut :: Pos -> Puzzle -> Int -> Int
numValuesRuledOut pos@(x, y) (Puzzle {..}) n =
    countTrue $ [puzzleRows `has` y' && puzzleBlocks `has` posBlock b'
                 | b'@(x', y') <- gridCols x, y /= y']
             ++ [puzzleCols `has` x' && puzzleBlocks `has` posBlock b'
                 | b'@(x', y') <- gridRows y, x /= x']
             ++ [puzzleCols `has` x' && puzzleRows `has` y'
                 | let b = posBlock pos, b'@(x', y') <- gridBlocks b, b /= b']
  where
    slots `has` ix = n `Set.member` find ix slots
    countTrue = foldl (\n b -> if b then n + 1 else n) 0

numConstraints pos@(x, y) (Puzzle {..}) =
    Set.size (find x puzzleCols) + Set.size (find y puzzleRows)
    + Set.size (find (posBlock pos) puzzleBlocks)

gridCols x = [(x, y) | y <- [0..8]]
gridRows y = [(x, y) | x <- [0..8]]
gridBlocks (x, y) = [(x', y') | x' <- [x*3..x*3+2], y' <- [y*3..y*3+2]]

makePuzzle ss =
    let ss'    = zip (filter (`elem` " 123456789") ss)
                     [(x, y) | y <- [0..8], x <- [0..8]]
        (m, o) = if length ss' /= 81 then error "Bad puzzle"
                 else foldl insertElem (Map.empty, Set.empty) ss'
    in Puzzle m o
       (Map.fromList [(x, removeNums m (gridCols x)) | x <- [0..8]])
       (Map.fromList [(y, removeNums m (gridRows y)) | y <- [0..8]])
       (Map.fromList [(pos, removeNums m (gridBlocks pos))
                      | x <- [0..2], y <- [0..2], let pos = (x, y)])
  where
    insertElem (m, o) (s, pos) | s == ' '  = (m, pos `Set.insert` o)
                               | otherwise = (Map.insert pos (ord s - ord '0')
                                                         m, o)

    removeNums m = foldl (\s -> maybe s (`Set.delete` s) . (`Map.lookup` m))
                         (Set.fromList [1..9])

puzzleStr (Puzzle {..}) =
    unlines [[maybe ' ' (chr . (ord '0' +)) ((x, y) `Map.lookup` puzzleGrid)
              | x <- [0..8]] | y <- [0..8]]

-- SOLVER

main = do [fc, h, level] <- getArgs
          let Just puzzle     = lookup level levels
              (solution, ctr) = backtrackingSearch (read fc) (read h) puzzle
          putStrLn $ puzzleStr puzzle
          putStrLn $ maybe "No solution." puzzleStr solution
          putStrLn $ "Searched " ++ show ctr ++ " nodes."

backtrackingSearch forwardChecking heuristics p = search p 0
  where
    search p n = case order [(mcv h, h) | pos <- Set.toList (puzzleOpen p),
                             let h = (pos, availableNumbers pos p)] of
                 ((_, pos):_) -> uncurry searchPos pos
                 []           -> (Just p, n)
      where
        mcv (pos, nums) = (Set.size nums, -numConstraints pos p)
        searchPos pos available = let vals = [(lcv val, val)
                                              | val <- Set.toList available]
            in foldl nextNode (Nothing, n) (snd `map` order vals)
          where
            lcv = numValuesRuledOut pos p
            nextNode (Nothing, n) val | discardEarly = (Nothing, n+1)
                                      | otherwise    = search p' (n+1)
              where
                discardEarly = forwardChecking && not (forwardCheck pos p')
                p'           = puzzleInsert pos val p
            nextNode solution _ = solution
        order :: Ord a => [a] -> [a]
        order = if heuristics then sort else id

easy = makePuzzle " 2 |6 5|   \
                  \ 75|   |  2\
                  \  9| 73|581\
                  \---+---+---\
                  \2  |81 |  9\
                  \ 8 |5 9| 4 \
                  \9  | 34|  8\
                  \---+---+---\
                  \491|38 |6  \
                  \5  |   |81 \
                  \   |4 1| 2 "

medium = makePuzzle " 82|46 |   \
                    \ 5 |73 | 2 \
                    \1  | 8 |   \
                    \---+---+---\
                    \2 4|   |1  \
                    \   |1 6|   \
                    \  6|   |3 8\
                    \---+---+---\
                    \   | 9 |  6\
                    \ 2 | 78| 1 \
                    \   | 14|73 "

hard = makePuzzle "  6|   |1  \
                  \7  | 61|   \
                  \ 1 |   |5  \
                  \---+---+---\
                  \ 9 |1 6|7  \
                  \5  |973|  8\
                  \  2|4 8| 9 \
                  \---+---+---\
                  \  5|   | 2 \
                  \   |83 |  6\
                  \  8|   |9  "

evil = makePuzzle " 2 |  8|6  \
                  \8  |  7|   \
                  \ 73| 4 |   \
                  \---+---+---\
                  \2  |  5| 9 \
                  \5  |6 9|  4\
                  \ 9 |3  |  5\
                  \---+---+---\
                  \   | 9 |31 \
                  \   |2  |  7\
                  \  2|4  | 8 "

levels = zip ["easy", "medium", "hard", "evil"] [easy, medium, hard, evil]