module Exercise03 where

import           Data.Array
import           Data.Maybe (catMaybes)
import           Data.IntSet (IntSet, difference, fromList, fromDistinctAscList, findMin)
import           Data.List   (intercalate, minimumBy, nub, sortBy)
import           Data.Ord    (comparing)
-- import           Debug.Trace (trace)
import           Text.Printf (printf)

import qualified Data.IntSet as IntSet

-- HA 3.1a) i
selectRow :: [[Int]] -> Int -> [Int]
selectRow xss i = xss !! i

-- HA 3.1a) ii
selectColumn :: [[Int]] -> Int -> [Int]
selectColumn xss i = [ xs !! i | xs <- xss ]

-- HA 3.1a) iii
intRoot :: Int -> Int
intRoot = floor . sqrt . fromIntegral

slice :: Int -> Int -> [a] -> [a]
slice i j xs = take (j - i) $ drop i xs

--return numbers in square as a list. squares are numbered  from left to right and top to bottom
--e.g. :
--[0,1,2]
--[3,4,5]
--[6,7,8]
selectSquare :: [[Int]] -> Int -> [Int]
selectSquare xss i = let r = intRoot $ length xss
                         x = r * rem i r
                         y = r * quot i r
                     in concatMap (slice x (x + r)) $ slice y (y + r) xss

-- HA 3.1b)
isValidSubsection :: [Int] -> Bool
isValidSubsection xs = let filled = filter (0 /=) xs
                       in filled == nub filled

isValidSudoku :: [[Int]] -> Bool
isValidSudoku []  = False
isValidSudoku xss = let ns = [0..length xss - 1]
                    in all isValidSubsection $ concat [xss, map (selectColumn xss) ns, map (selectSquare xss) ns]

-- HA 3.1c)
setCell :: [[Int]] -> (Int,Int) -> Int -> [[Int]]
setCell xss (j, k) x = let (r0, r:r1) = splitAt j xss
                           (c0, c:c1) = splitAt k r
                        in r0 ++ (c0 ++ x:c1):r1

-- HA 3.1d)
{-WETT-}

-- Alias for zip [0..] xs, helpful when dealing with indexed values
enumerate :: [a] -> [(Int,a)]
enumerate = zip [0..]

-- Packing the important information in a special data structure to reduce time evaluating trunks
data Cell = Fixed Int | Possible IntSet deriving (Show)
data Sudoku = Sudoku { board :: Array Int Cell, size :: Int, root :: Int } deriving (Show)

-- Pack the Sudoku in the custom Data Structure
buildSudoku :: [[Int]] -> Sudoku
buildSudoku xss = let n = length xss
                      r = intRoot n
                  in Sudoku { board=buildCells xss n r, size=n, root=r }

-- Builds an array of cells from the Sudoku
buildCells :: [[Int]] -> Int -> Int -> Array Int Cell
buildCells xss n r = listArray (0, (n * n) - 1) (buildCells' xss n r)

-- Builds a list of cells from the Sudoku
buildCells' :: [[Int]] -> Int -> Int -> [Cell]
buildCells' xss n r = let aux 0 xs = Possible xs
                          aux x _  = Fixed x
                      in [ aux x (validMoves xss n r i) | (i, x) <- zip [0..] $ concat xss ]

-- Return the valid moves in a certain cell
validMoves :: [[Int]] -> Int -> Int -> Int -> IntSet
validMoves xss n r i = let (y, x) = quotRem i n
                           k      = quot x r + r * quot y r
                           set    = fromList $ selectRow xss y ++ selectColumn xss x ++ selectSquare xss k
                       in difference (fromDistinctAscList [1..n]) set

-- Get a row of a Sudoku
getRow :: Sudoku -> Int -> Array Int Cell
getRow sudoku r = let b = board sudoku
                      n = size sudoku
                  in ixmap (0, n - 1) (\i -> i + n * r) b

-- Get a column of a Sudoku
getColumn :: Sudoku -> Int -> Array Int Cell
getColumn sudoku c = let b = board sudoku
                         n = size sudoku
                     in ixmap (0, n - 1) (\i -> n * i + c) b

-- Get a sub square of a Sudoku
getSquare :: Sudoku -> Int -> Array Int Cell
getSquare sudoku s = let b      = board sudoku
                         n      = size sudoku
                         r      = root sudoku
                         (y, x) = quotRem s r
                         ixs i  = let (y', x') = quotRem i r in (y * n + x) * r + x' + y' * n
                     in ixmap (0, n - 1) ixs b

-- Get the indices of a Row
getRowIxs :: Sudoku -> Int -> [Int]
getRowIxs ss j = let n = size ss in [j * n..(j + 1) * n - 1]

-- Get the indices of a Column
getColumnIxs :: Sudoku -> Int -> [Int]
getColumnIxs ss i = let n = size ss in map (\ix -> i + ix * n) [0..n - 1]

-- Get the indices of a Sub Square
getSquareIxs :: Sudoku -> Int -> Int -> [Int]
getSquareIxs ss i j = let n = size ss
                          r = root ss
                          (x, y) = (r * quot i r, r * quot j r)
                          offset = y * n + x
                      in map (\ix -> let (y', x') = quotRem ix r in offset + x' + y' * n) [0..n - 1]

-- Update the board of a Sudoku
updateBoard :: Sudoku -> Array Int Cell -> Sudoku
updateBoard ss b = Sudoku {board=b, size=size ss, root=root ss}

-- Update a Cell's Possibilities
updateCell :: Cell -> Int -> Cell
updateCell (Possible xs) x = Possible (IntSet.delete x xs)
updateCell c _ = c

-- Get Cells that need to be updated after changing a Cell
-- Sudoku -> Index -> [Cell Indices]
getDirtyCells :: Sudoku -> Int -> [Int]
getDirtyCells ss ix = let (y, x) = quotRem ix n 
                          n      = size ss
                      in nub (getRowIxs ss y ++ getColumnIxs ss x ++ getSquareIxs ss x y)

-- Update afected Cell possibilities in the Sudoku
-- Sudoku -> [(Cell Index, Cell Value)] -> Sudoku
updatePossible :: Sudoku -> [(Int, Int)] -> Sudoku
updatePossible ss []         = ss
updatePossible ss ((i,x):xs) = let b  = board ss
                                   b' = b // [(ix, updateCell (b ! ix) x) | ix <- getDirtyCells ss i]
                               in updatePossible (updateBoard ss b') xs

-- Set all Cells that only have one possibility
pruneSudoku :: Sudoku -> Maybe Sudoku
pruneSudoku ss = let b  = board ss
                     is = [(i,Fixed (findMin xs)) | (i,Possible xs) <- enumerate (elems b),IntSet.size xs == 1]
                     s' =  updatePossible (updateBoard ss (b // is)) $ map (\(i, Fixed x) -> (i, x)) is
                 in if null is then validateSudoku ss else pruneSudoku s'

setCells :: Sudoku -> [(Int, Int)]
setCells ss = let b = board ss in [ (i, x) | (i, Fixed x) <- enumerate $ elems b ]

unsetCells :: Sudoku -> [(Int, IntSet)]
unsetCells ss = let b = board ss in [ (i, x) | (i, Possible x) <- enumerate $ elems b ]

fixCell :: Sudoku -> Int -> Int -> Maybe Sudoku
fixCell ss i x = let ss' = updatePossible (updateBoard ss $ board ss // [(i, Fixed x)]) [(i, x)]
                 in validateSudoku ss'

validateSudoku :: Sudoku -> Maybe Sudoku
validateSudoku ss = if invalidSudoku ss then Nothing else Just ss

-- True if the Sudoku is Solved
solvedSudoku :: Sudoku -> Bool
solvedSudoku ss = let b = board ss in not $ null $ unsetCells ss

-- 
validFixedCells :: Sudoku -> Bool
validFixedCells ss = let r = [ [ x | (Fixed x) <- elems $ getRow ss i ] | i <- [0..size ss - 1] ]
                         c = [ [ x | (Fixed x) <- elems $ getColumn ss i ] | i <- [0..size ss - 1] ]
                         s = [ [ x | (Fixed x) <- elems $ getSquare ss i ] | i <- [0..size ss - 1] ]
                         valid [] = True
                         valid (xs:xss) = length (nub xs) == length xs && valid xss
                     in valid r && valid c && valid s

-- True if any Cell has no Possibilities and isn't Fixed
invalidSudoku :: Sudoku -> Bool
invalidSudoku ss = let b = board ss
                       n = any (IntSet.null . snd) $ unsetCells ss
                       v = not $ validFixedCells ss
                   in n || v 

-- A head function that doesn't crash when given the empty list
safeHead :: [a] -> Maybe a
safeHead []     = Nothing
safeHead (x:xs) = Just x

-- Solve Sudoku
solveSudoku :: [[Int]] -> [[Int]]
solveSudoku xss = unpackSudoku $ solveSudoku' $ buildSudoku xss

solveSudoku' :: Sudoku -> Maybe Sudoku
solveSudoku' ss = do s' <- pruneSudoku ss
                     helper s' $ unsetCells s'

helper :: Sudoku -> [(Int, IntSet)] -> Maybe Sudoku                     
helper ss us = let (i, xs) = head us
               in if null us then Just ss 
                  else safeHead $ catMaybes [fixCell ss i x >>= solveSudoku' | x <- IntSet.elems xs ]

-- Transform Custom Sudoku into Exercise Sudoku
unpackSudoku :: Maybe Sudoku -> [[Int]]
unpackSudoku Nothing   = []
unpackSudoku (Just ss) = let n               = size ss
                             b               = board ss
                             f (Fixed x)     = x
                             f (Possible xs) = 0
                         in [[f (b ! (x + y * n)) | x <- [0..n-1]] | y <- [0..n-1]]

{-TTEW-}

easySudoku :: [[Int]]
easySudoku = [[0,0,1,0],
              [0,3,2,0],
              [0,2,3,0],
              [0,1,0,0]]

mediSudoku :: [[Int]]
mediSudoku = [[0,5,8,0,7,0,0,0,2],
              [0,4,0,0,6,2,0,9,8],
              [2,9,1,0,3,0,7,0,0],
              [0,0,6,9,0,0,4,0,7],
              [3,2,0,6,0,0,0,1,5],
              [0,7,0,2,5,4,6,0,3],
              [0,0,0,8,9,1,2,0,6],
              [0,0,0,0,2,0,0,4,0],
              [0,0,0,0,0,0,8,0,1]]

hardSudoku :: [[Int]]
hardSudoku = [[8,0,0,0,0,0,0,0,0],
              [0,0,3,6,0,0,0,0,0],
              [0,7,0,0,9,0,2,0,0],
              [0,5,0,0,0,7,0,0,0],
              [0,0,0,0,4,5,7,0,0],
              [0,0,0,1,0,0,0,3,0],
              [0,0,1,0,0,0,0,6,8],
              [0,0,8,5,0,0,0,1,0],
              [0,9,0,0,0,0,4,0,0]]

-- Utility method to show a sudoku
-- show sudoku with
-- >>> putStr (showSudoku sudoku)
showSudoku :: [[Int]] -> String
showSudoku xss = unlines $ intercalate [showDivider] $ chunksOf squareSize $ map showRow xss
  where
    size = length xss
    squareSize = intRoot size
    numberSize = size `div` 10 + 1

    showRowSection xs = unwords $ map (printf ("%0" ++ show numberSize ++ "d")) xs
    showRow xs = intercalate "|" $ map showRowSection $ chunksOf squareSize xs
    showDivider = intercalate "+" $ replicate squareSize $ replicate ((numberSize + 1) * squareSize - 1) '-'

    chunksOf :: Int -> [e] -> [[e]]
    chunksOf i [] = []
    chunksOf i ls = take i ls : chunksOf i (drop i ls)