hasdoku/src/Sudoku.hs

module Main where

import Control.Applicative ((<|>))
import Data.Function ((&))
import Data.Vector ((!))
import qualified Control.Monad
import qualified Control.Monad.ST
import qualified Data.Char
import qualified Data.Function
import qualified Data.List
import qualified Data.Map.Strict as Map
import qualified Data.Word
import qualified Data.Bits
import qualified Data.Vector
import qualified Data.Vector.Unboxed
import qualified Data.Vector.Unboxed.Mutable
import qualified Data.STRef

fixM :: (Eq t, Monad m) => (t -> m t) -> t -> m t
fixM f x = f x >>= \x' -> if x' == x then return x else fixM f x'

setBits :: Data.Word.Word16 -> [Data.Word.Word16] -> Data.Word.Word16
setBits = Data.List.foldl' (Data.Bits..|.)

data Cell = Fixed Data.Word.Word16
          | Possible Data.Word.Word16
          deriving (Show, Eq)

type Grid = Data.Vector.Vector Cell
type CellIxs = [Int]

isPossible :: Cell -> Bool
isPossible (Possible _) = True
isPossible _            = False

readGrid :: String -> Maybe Grid
readGrid s
  | length s == 81 = Data.Vector.fromList <$> traverse readCell s
  | otherwise      = Nothing
  where
    allBitsSet = 1022

    readCell '.' = Just $ Possible allBitsSet
    readCell c
      | Data.Char.isDigit c && c > '0' = Just . Fixed . Data.Bits.bit . Data.Char.digitToInt $ c
      | otherwise = Nothing

fromXY :: (Int, Int) -> Int
fromXY (x, y) = x * 9 + y

allRowIxs, allColIxs, allSubGridIxs :: [CellIxs]
allRowIxs = [getRow i | i <- [0..8]]
  where getRow n = [ fromXY (n, i) | i <- [0..8] ]

allColIxs = [getCol i | i <- [0..8]]
  where getCol n = [ fromXY (i, n) | i <- [0..8] ]

allSubGridIxs = [getSubGrid i | i <- [0..8]]
  where getSubGrid n = let (r, c) = (n `quot` 3, n `mod` 3)
                       in [ fromXY (3 * r + i, 3 * c + j) | i <- [0..2], j <- [0..2] ]

replaceCell :: Int -> Cell -> Grid -> Grid
replaceCell i c g = g Data.Vector.// [(i, c)]

showGrid :: Grid -> String
showGrid grid = unlines . map (unwords . map (showCell . (grid !))) $ allRowIxs
  where
    showCell (Fixed x) = show . Data.Bits.countTrailingZeros $ x
    showCell _         = "."

showGridWithPossibilities :: Grid -> String
showGridWithPossibilities grid = unlines . map (unwords . map (showCell . (grid !))) $ allRowIxs
  where
    showCell (Fixed x)     = (show . Data.Bits.countTrailingZeros $ x) ++ "          "
    showCell (Possible xs) =
      "[" ++ map (\i -> if Data.Bits.testBit xs i then Data.Char.intToDigit i else ' ') [1..9] ++ "]"

cellIndicesList :: [Cell] -> [Data.Word.Word16]
cellIndicesList cells =
  Data.Vector.Unboxed.toList $ Control.Monad.ST.runST $ do
    vec <- Data.Vector.Unboxed.Mutable.replicate 9 Data.Bits.zeroBits
    ref <- Data.STRef.newSTRef (1 :: Int)
    Control.Monad.forM_ cells $ \cell -> do
      i <- Data.STRef.readSTRef ref
      case cell of
        Fixed _ -> return ()
        Possible xs -> Control.Monad.forM_ [0..8] $ \d ->
          Control.Monad.when (Data.Bits.testBit xs (d+1)) $
            Data.Vector.Unboxed.Mutable.unsafeModify vec (`Data.Bits.setBit` i) d
      Data.STRef.writeSTRef ref (i+1)
    Data.Vector.Unboxed.unsafeFreeze vec

exclusivePossibilities :: [Cell] -> [Data.Word.Word16]
exclusivePossibilities row =
  row
  & cellIndicesList
  & zip [1..9]
  & filter (\(_, xs) -> let p = Data.Bits.popCount xs in p > 0 && p < 4)
  & Data.List.foldl' (\acc (x, is) -> Map.insertWith prepend is [x] acc) Map.empty
  & Map.filterWithKey (\is xs -> Data.Bits.popCount is == length xs)
  & Map.elems
  & map (Data.List.foldl' Data.Bits.setBit Data.Bits.zeroBits)
  where
    prepend ~[y] ys = y:ys

makeCell :: Data.Word.Word16 -> Maybe Cell
makeCell ys
  | ys == Data.Bits.zeroBits   = Nothing
  | Data.Bits.popCount ys == 1 = Just $ Fixed ys
  | otherwise                  = Just $ Possible ys

pruneCells :: Grid -> CellIxs -> Maybe Grid
pruneCells grid cellIxs = Control.Monad.foldM pruneCell grid cellIxs
  where
    cells         = map (grid !) cellIxs
    exclusives    = exclusivePossibilities cells
    allExclusives = setBits Data.Bits.zeroBits exclusives
    fixeds        = setBits Data.Bits.zeroBits [x | Fixed x <- cells]

    pruneCell g i =
      pruneCellByFixed g (i, g ! i) >>= \g' -> pruneCellByExclusives g' (i, g' ! i)

    pruneCellByFixed g (_, Fixed _) = Just g
    pruneCellByFixed g (i, Possible xs)
      | xs' == xs = Just g
      | otherwise = flip (replaceCell i) g <$> makeCell xs'
      where
        xs' = xs Data.Bits..&. Data.Bits.complement fixeds

    pruneCellByExclusives g (_, Fixed _) = Just g
    pruneCellByExclusives g (i, Possible xs)
      | null exclusives                = Just g
      | intersection == xs             = Just g
      | intersection `elem` exclusives =
          flip (replaceCell i) g <$> makeCell intersection
      | otherwise                      = Just g
      where
        intersection = xs Data.Bits..&. allExclusives

pruneGrid' :: Grid -> Maybe Grid
pruneGrid' grid =
  Control.Monad.foldM pruneCells grid allRowIxs
  >>= flip (Control.Monad.foldM pruneCells) allColIxs
  >>= flip (Control.Monad.foldM pruneCells) allSubGridIxs

pruneGrid :: Grid -> Maybe Grid
pruneGrid = fixM pruneGrid'

isGridFilled :: Grid -> Bool
isGridFilled = not . Data.Vector.any isPossible

isGridInvalid :: Grid -> Bool
isGridInvalid grid =
  any isInvalidRow (map (map (grid !)) allRowIxs)
  || any isInvalidRow (map (map (grid !)) allColIxs)
  || any isInvalidRow (map (map (grid !)) allSubGridIxs)
  where
    isInvalidRow row =
      let fixeds         = [x  | Fixed x <- row]
          emptyPossibles = [() | Possible x <- row, x == Data.Bits.zeroBits]
      in hasDups fixeds || not (null emptyPossibles)

    hasDups l = hasDups' l []

    hasDups' [] _ = False
    hasDups' (y:ys) xs
      | y `elem` xs = True
      | otherwise   = hasDups' ys (y:xs)

nextGrids :: Grid -> (Grid, Grid)
nextGrids grid =
  let (i, first@(Fixed _), rest) =
        fixCell
        . Data.Vector.minimumBy (compare `Data.Function.on` (possibilityCount . snd))
        . Data.Vector.imapMaybe (\j cell -> if isPossible cell then Just (j, cell) else Nothing)
        $ grid
  in (replaceCell i first grid, replaceCell i rest grid)
  where
    possibilityCount (Possible xs) = Data.Bits.popCount xs
    possibilityCount (Fixed _)     = 1

    fixCell ~(i, Possible xs) =
      let x = Data.Bits.countTrailingZeros xs
      in case makeCell (Data.Bits.clearBit xs x) of
        Nothing -> error "Impossible case"
        Just cell -> (i, Fixed (Data.Bits.bit x), cell)

solve :: Grid -> Maybe Grid
solve grid = pruneGrid grid >>= solve'
  where
    solve' g
      | isGridInvalid g = Nothing
      | isGridFilled g  = Just g
      | otherwise       =
          let (grid1, grid2) = nextGrids g
          in solve grid1 <|> solve grid2

main :: IO ()
main = do
  inputs <- lines <$> getContents
  Control.Monad.forM_ inputs $ \input ->
    case readGrid input of
      Nothing   -> putStrLn "Invalid input"
      Just grid -> case solve grid of
        Nothing    -> putStrLn "No solution found"
        Just grid' -> putStrLn $ showGrid grid'