module Main(main) where
import Data.Array
import Data.IntSet (IntSet)
import qualified Data.IntSet as IS
import Control.Monad ((>=>))
import Control.Applicative ((<|>))
import Data.Char
-- import Debug.Trace
import Data.List (foldl')
import Data.Maybe(fromMaybe)

data Cell = Fixed Int | Possible IntSet
  deriving (Show, Eq)
type Grid = Array (Int,Int) Cell

cellToChar :: Cell -> Char
cellToChar (Fixed n)    = intToDigit n
cellToChar (Possible _) = '.'

showCell :: Cell -> String
showCell (Fixed x)     = show x ++ "          "
showCell (Possible is) = (++ "]") . foldl' (\acc x -> acc ++ if x `IS.member` is then show x else " ") "[" $ [1..9]

showGrid :: Grid -> String
showGrid = unlines . fmap (fmap cellToChar) . chunksOf 9 . elems

showGridWithPossibilities :: Grid -> String
showGridWithPossibilities = unlines . map (unwords . map showCell) . rows
-- I should write index functions for rows and columns
-- then I could make an abstract access functions and define
-- the concrete ones using the index function and the abstract
-- accessors (same for the update ones)
-- (like prune pruneStructures)
row :: Grid -> Int -> [Cell]
row grid rn = [grid ! (rn,i) | i <- [0..8]]

rows :: Grid -> [[Cell]]
rows grid = [row grid i | i <- [0..8]]

column :: Grid -> Int -> [Cell]
column grid rn = [grid ! (i,rn) | i <- [0..8]]

columns :: Grid -> [[Cell]]
columns grid = [column grid i | i <- [0..8]]

sGrid :: Grid -> Int -> [Cell]
sGrid grid sGridn = [grid ! (a,b) | (a,b) <- intToSGindex sGridn]

sGrids :: Grid -> [[Cell]]
sGrids grid = [sGrid grid i | i <- [0..8]]

{-# INLINE intToSGindex #-}
intToSGindex :: Int -> [(Int,Int)]
intToSGindex 0 = [(x,y) | x <- [0..2], y <- [0..2] ]
intToSGindex 1 = [(x,y) | x <- [0..2], y <- [3..5] ]
intToSGindex 2 = [(x,y) | x <- [0..2], y <- [6..8] ]
intToSGindex 3 = [(x,y) | x <- [3..5], y <- [0..2] ]
intToSGindex 4 = [(x,y) | x <- [3..5], y <- [3..5] ]
intToSGindex 5 = [(x,y) | x <- [3..5], y <- [6..8] ]
intToSGindex 6 = [(x,y) | x <- [6..8], y <- [0..2] ]
intToSGindex 7 = [(x,y) | x <- [6..8], y <- [3..5] ]
intToSGindex 8 = [(x,y) | x <- [6..8], y <- [6..8] ]
intToSGindex _ = []

pruneCells :: [Cell] -> Maybe [Cell]
pruneCells cells = traverse pruneCell cells
  where
    fixeds = IS.fromList [x | Fixed x <- cells]

    pruneCell (Possible xs) = t $ xs IS.\\ fixeds
    pruneCell (Fixed x)     = Just $ Fixed x
    t x | IS.size x == 0 = Nothing
    t x | IS.size x == 1 = Just $ Fixed $ head $ IS.toList x
    t x | otherwise      = Just $ Possible x

checkConstraint :: Grid -> Maybe Grid
checkConstraint grid = case and [abs (x - y ) > 1 | i <- [0..7], j <- [0..8], Fixed x <- [grid ! (i,j)], Fixed y <- [grid ! (i+1,j)]] && and [abs (x - y) > 1 | i <- [0..8], j <- [0..7], Fixed x <- [grid ! (i,j)], Fixed y <- [grid ! (i,j+1)]] of
  True  -> Just grid
  False -> Nothing

updateRow :: Grid -> [Cell] -> Int -> Grid
updateRow grid cells rn = grid // zip [(rn,i) | i <- [0..8]] cells

updateRows :: Grid -> [[Cell]] -> Grid
updateRows grid cellss = grid // concat (zipWith f [0..8] cellss)
  where
    f i = zip [(i,x)| x <- [0..8]]

updateColumn :: Grid -> [Cell] -> Int -> Grid
updateColumn grid cells cn = grid // zip [(i,cn) | i <- [0..8]] cells

updateColumns :: Grid -> [[Cell]] -> Grid
updateColumns grid cells = grid // concat  (zipWith  f [0..8] cells)
  where
    f i = zip [(x,i)| x <- [0..8]]

updateSGrid :: Grid -> [Cell] -> Int -> Grid
updateSGrid grid cells sGn = grid // zip (intToSGindex sGn) cells

updateSGrids :: Grid -> [[Cell]] -> Grid
updateSGrids grid cells = grid  // concat (zipWith f [0..8] cells)
  where
    f i = zip (intToSGindex i)

{-# INLINE pruneStructures #-}
pruneStructures :: (Grid -> [[Cell]]) -> (Grid -> [[Cell]] -> Grid) -> Grid -> Maybe Grid
pruneStructures access update grid = do
  let rs = access grid
  prs <- traverse pruneCells rs
  return $ update grid prs

pruneRows :: Grid -> Maybe Grid
pruneRows = pruneStructures rows updateRows

pruneColummns :: Grid -> Maybe Grid
pruneColummns = pruneStructures columns updateColumns

pruneSGrids :: Grid -> Maybe Grid
pruneSGrids = pruneStructures sGrids updateSGrids

pruneGrid :: Grid -> Maybe Grid
pruneGrid = pruneRows >=> pruneColummns >=> pruneSGrids

fixGridMaybe :: (Grid -> Maybe Grid) -> Grid -> Maybe Grid
fixGridMaybe f grid = do
  x <- f grid
  if x == grid
    then return grid
    else fixGridMaybe f x

prunesGrid :: Grid -> Maybe Grid
prunesGrid = fixGridMaybe pruneGrid

isFixed :: Cell -> Bool
isFixed (Fixed _) = True
isFixed _         = False

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

allFixed :: Grid -> Bool
allFixed grid = all isFixed $ elems grid

headL :: [a] -> [a]
headL []    = []
headL (x:_) = [x]

travTuple :: (Maybe a, Maybe b) -> Maybe (a,b)
travTuple (Nothing, _)     = Nothing
travTuple (_,Nothing)      = Nothing
travTuple (Just x, Just y) = Just (x,y)

makeChoice :: Grid -> (Grid,Grid)
makeChoice grid = splitGrid
  where
    smallestCellInd :: ((Int,Int),[Int])
    smallestCellInd = convert . indHelper $ (assocs grid)
    convert (c,Possible is) = (c,IS.toList is)
    convert (_,Fixed _)     = error "convert error"
    indHelper []                      = error "We have no choices but aren't done?"
    indHelper ((_,Fixed _) : xs)      = indHelper xs
    indHelper (t@(_,Possible _) : xs) = indHelper' t xs
    indHelper' :: ((Int,Int),Cell) -> [((Int,Int),Cell)] -> ((Int,Int),Cell)
    indHelper' t [] = t
    indHelper' p@(_,Possible is) (p'@(_, Possible is') : xs) = if IS.size is <= IS.size is' then indHelper' p xs else indHelper' p' xs
    indHelper' p ((_,(Fixed _)) : xs)                        = indHelper' p xs
    indHelper' (_,Fixed _) _                                 = error "indHelper' impossible branch"
    splitGrid :: (Grid,Grid)
    splitGrid = case smallestCellInd of
      (c,[x,y])     -> (grid // [(c,Fixed x)], grid // [(c,Fixed y)])
      (c,x:t@(_:_)) -> (grid // [(c,Fixed x)], grid // [(c,Possible $ IS.fromList t)])
      _             -> error "We just checked the size"
        
solveSodoku :: Grid -> Maybe Grid
solveSodoku grid = do
  nGrid <- prunesGrid grid
  nGrid' <- checkConstraint nGrid
  if allFixed nGrid'
    then return nGrid
    else do
      let (x,y) = makeChoice nGrid
      solveSodoku x <|> solveSodoku y

readSodoku :: String -> Maybe Grid
readSodoku str | length str == 81 = fmap (listArray ((0,0),(8,8))) . traverse readCell $ str
               | otherwise        = Nothing
  where
    readCell :: Char -> Maybe Cell
    readCell '.' = Just . Possible . IS.fromList $ [1..9]
    readCell c | isDigit c && c > '0' = Just . Fixed . digitToInt $ c
               | otherwise            = Nothing


chunksOf :: Int -> [a] -> [[a]]
chunksOf _ [] = []
chunksOf n xs = let (c,txs) = splitAt n xs in c : chunksOf n txs

{-# NOINLINE sodokuStr #-}
sodokuStr :: String
sodokuStr = "....4.....9.....6...4.9.1..............1.3..............3.6.2...2.....4.....7...."

main :: IO ()
main = do
  putStrLn $ fromMaybe "Unsolvable" $ do
    sodoku <- readSodoku sodokuStr
    solution <- solveSodoku sodoku
    return $ showGrid solution