-- Based on the paper "A Program to Solve Sudoku" by Richard Bird.
-- Written by Ziv for a SIPB IAP class in January 2015.
-- Feel free to email learn-haskell@mit.edu or ziv@mit.edu with questions!

module Sudoku where

import Data.List (transpose, nub, (\\), intersperse)

-----------
-- specs --
-----------
type Entry = Char

size :: Int
size = 3

allNumbers :: [Entry]
allNumbers = "123456789"

empty :: Entry -> Bool
empty x = x == '.'

-------------
-- utility --
-------------
clump :: Int -> [a] -> [[a]]
clump _ [] = []
clump n xs = take n xs : clump n (drop n xs)

(~>) :: (a -> b) -> (b -> c) -> (a -> c)
f ~> g = g . f

data Choice a = OneOf [a]
    deriving (Eq, Show)

choices :: Choice a -> [a]
choices (OneOf xs) = xs

noDups :: Eq a => [a] -> Bool
noDups xs = xs == nub xs

lcp :: [Choice a] -> Choice [a]
lcp = map choices ~> sequence ~> OneOf
-- *handwave*

mcp :: [[Choice a]] -> Choice [[a]]
mcp = map lcp ~> lcp

oneChoice :: Choice a -> Bool
oneChoice (OneOf [_]) = True
oneChoice _ = False

repeatUntilFutile :: Eq a => (a -> a) -> a -> a
repeatUntilFutile f x = let y = f x
                        in if x == y then x else repeatUntilFutile f y

concatChoice :: [Choice a] -> Choice a
concatChoice = map choices ~> concat ~> OneOf

-----------
-- basic --
-----------
rows :: [[a]] -> [[a]]
rows xss = xss

cols :: [[a]] -> [[a]]
cols = transpose

boxs :: [[a]] -> [[a]]
boxs = map (clump size) ~> clump size ~> map transpose ~> concat ~> map concat

choiceEntry :: Entry -> Choice Entry
choiceEntry x = if empty x then OneOf allNumbers else OneOf [x]

choiceMatrix :: [[Entry]] -> [[Choice Entry]]
choiceMatrix = map (map choiceEntry)

validBy :: Eq a => ([[a]] -> [[a]]) -> [[a]] -> Bool
validBy f = f ~> all noDups

valid :: Eq a => [[a]] -> Bool
valid xss = validBy rows xss && validBy cols xss && validBy boxs xss

solveBasic :: [[Entry]] -> [[Entry]]
solveBasic = choiceMatrix ~> mcp ~> choices ~> filter valid ~> head

-------------
-- Pruning --
-------------
pruneList :: Eq a => [Choice a] -> [Choice a]
pruneList xs = let fixeds = concat (map choices (filter oneChoice xs))
                   pruneElem x = if oneChoice x
                                 then x
                                 else OneOf (choices x \\ fixeds)
               in map pruneElem xs

pruneBy :: Eq a => ([[Choice a]] -> [[Choice a]]) -> [[Choice a]] -> [[Choice a]]
pruneBy f = f ~> map pruneList ~> f

prune :: Eq a => [[Choice a]] -> [[Choice a]]
prune = repeatUntilFutile (pruneBy rows ~> pruneBy cols ~> pruneBy boxs)

{-
  mcp ~> choices ~> filter valid
    = prune ~> mcp ~> choices ~> filter valid
 -}

solvePrune :: [[Entry]] -> [[Entry]]
solvePrune = choiceMatrix ~> prune ~> mcp ~> choices ~> filter valid ~> head

--------------
-- Guessing --
--------------
guessList :: [Choice a] -> [[Choice a]]
guessList (x : xs) = if oneChoice x
                     then map (\ys -> x : ys) (guessList xs)
                     else map (\c -> OneOf [c] : xs) (choices x)

guess :: [[Choice a]] -> [[[Choice a]]]
guess = concat ~> guessList ~> map (clump (size^2))

{-
0  mcp = guess ~> map mcp ~> concatChoice
1  mcp ~> choices ~> filter valid
2    = guess ~> map mcp ~> concatChoice ~> choices ~> filter valid
3    = guess ~> map mcp ~> choices ~> concat ~> filter valid
4    = guess ~> map (mcp ~> choices ~> filter valid) ~> concat
5    = guess ~> map (prune ~> mcp ~> choices ~> filter valid) ~> concat
 -}

search :: Eq a => [[Choice a]] -> [[[a]]]
search xss = if all (all oneChoice) xss
             then if valid xss then choices (mcp xss) else []
             else (guess ~> map (prune ~> search) ~> concat) xss

solve :: [[Entry]] -> [[Entry]]
solve = choiceMatrix ~> prune ~> search ~> head

-------------
-- Testing --
-------------
grid = ["39...2..6",
        ".5..86...",
        "2.......3",
        ".3.7.....",
        "..1.6.8..",
        ".....1.9.",
        "4.......7",
        "...43..5.",
        "...6...32"]

test1 = ["1234","4321","3412","2143"]
test2 = ["1.34","43.1","3412","21.3"]
test3 = [".2..",".3.1","3...","2..3"]
test4 = [".2..",".3.1","3...","2..."]

printGrid = putStrLn . unlines . map (intersperse ' ')