#!/usr/bin/runghc

module Recursion14Sudoku where

import Data.List as L
import Data.Vector as V

exampleGrid :: Vector (Vector Int)
exampleGrid =
  V.fromList
    ( Prelude.map
        V.fromList
        [ [3, 0, 6, 5, 0, 8, 4, 0, 0],
          [5, 2, 0, 0, 0, 0, 0, 0, 0],
          [0, 8, 7, 0, 0, 0, 0, 3, 1],
          [0, 0, 3, 0, 1, 0, 0, 8, 0],
          [9, 0, 0, 8, 6, 3, 0, 0, 5],
          [0, 5, 0, 0, 9, 0, 6, 0, 0],
          [1, 3, 0, 0, 0, 0, 2, 5, 0],
          [0, 0, 0, 0, 0, 0, 0, 7, 4],
          [0, 0, 5, 2, 0, 6, 3, 0, 0]
        ]
    )

tinyGrid :: Vector (Vector Int)
tinyGrid =
  V.fromList
    ( Prelude.map
        V.fromList
        [ [1, 2],
          [3, 4]
        ]
    )

-- |
-- >>> renderGrid tinyGrid
-- "[1,2]\n[3,4]"
renderGrid :: Vector (Vector Int) -> String
renderGrid grid = L.intercalate "\n" (toList (V.map show grid))

-- |
-- >>> inRow exampleGrid 0 6
-- True

-- |
-- >>> inRow exampleGrid 0 7
-- False
inRow :: Vector (Vector Int) -> Int -> Int -> Bool
inRow exampleGrid row num = V.elem num (exampleGrid V.! row)

-- |
-- >>> inCol exampleGrid 1 8
-- True

-- |
-- >>> inCol exampleGrid 1 6
-- False
inCol :: Vector (Vector Int) -> Int -> Int -> Bool
inCol exampleGrid col num = V.elem num (V.map (V.! col) exampleGrid)

-- |
-- >>> inBox exampleGrid 0 0 2
-- True

-- |
-- >>> inBox exampleGrid 0 0 11
-- False
inBox :: Vector (Vector Int) -> Int -> Int -> Int -> Bool
inBox grid startRow startCol num =
  let box = V.map (V.slice startCol 3) (V.slice startRow 3 grid)
   in V.any id (V.map (V.elem num) box)

-- |
-- >>> isSafe exampleGrid 0 0 2
-- False

-- |
-- >>> isSafe exampleGrid 0 0 11
-- True
isSafe :: Vector (Vector Int) -> Int -> Int -> Int -> Bool
isSafe grid row col num =
  not (inRow grid row num) && not (inCol grid col num) && not (inBox grid (row - mod row 3) (col - mod col 3) num)

-- |
-- >>> firstEmpty exampleGrid 0
-- Just (0,1)

-- |
-- >>> firstEmpty exampleGrid 1
-- Just (1,2)
firstEmpty :: Vector (Vector Int) -> Int -> Maybe (Int, Int)
firstEmpty grid row = case grid V.!? row of
  Nothing -> Nothing
  Just currRow -> case V.elemIndex 0 currRow of
    Just col -> Just (row, col)
    Nothing -> firstEmpty grid (row + 1)

solveSudoku :: Vector (Vector Int) -> Maybe (Vector (Vector Int))
solveSudoku grid = case firstEmpty grid 0 of
  Nothing -> Just grid
  Just (row, col) ->
    let moveCheck num =
          if 9 < num
            then Nothing
            else
              if isSafe grid row col num
                then
                  let newRow = (grid V.! row) // [(col, num)]
                      newGrid = grid // [(row, newRow)]
                   in case solveSudoku newGrid of
                        Just solvedGrid -> Just solvedGrid
                        Nothing -> moveCheck (num + 1)
                else moveCheck (num + 1)
     in moveCheck 1

main :: IO ()
main = do
  putStrLn (renderGrid exampleGrid)
  putStrLn "\nSolving board\n"
  putStrLn
    ( case solveSudoku exampleGrid of
        Just grid -> renderGrid grid
        Nothing -> "Solution not found"
    )