module Exercise03 where

import Text.Printf (printf)
import Data.List as List
import Data.IntSet as IntSet
import Data.Set as Set

-- 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

--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 = concat [ List.take num $ List.drop col $ selectRow xss j | j <- [row .. row + num - 1] ]
    where
        num = intRoot (length xss) -- number of squares per row / column
        row = (i `div` num) * num -- get index of first row of square
        col = (i `mod` num) * num -- get index of first column of square

-- HA 3.1b)
isValidSubsection :: [Int] -> Bool
isValidSubsection [] = True
isValidSubsection (x:xs) = (x == 0 || x `notElem` xs) && isValidSubsection xs

isValidSudoku :: [[Int]] -> Bool
isValidSudoku xss = validRows && validColumns && validSquares
    where
        range = [0..length xss - 1]
        validColumns = all isValidSubsection [ selectColumn xss i | i <- range ]
        validRows = all isValidSubsection [ selectRow xss i | i <- range ]
        validSquares = all isValidSubsection [ selectSquare xss i | i <- range ]

-- HA 3.1c)
setCell :: [[Int]] -> (Int,Int) -> Int -> [[Int]]
setCell xss (r, c) x = List.take r xss ++ [newRow] ++ List.drop (r + 1) xss
    where
        row = xss !! r
        newRow = List.take c row ++ [x] ++ List.drop (c + 1) row

-- HA 3.1d)
{-WETT-}

solveSudoku :: [[Int]] -> [[Int]]
solveSudoku xss =
    if isValidSudoku xss then
        solveSudokuIntern xss possibilitiesMap
    else
        []
    where
        len = length xss
        possibilitiesMap = [
            ((x,y),possibilitiesForIndex xss (x,y))
            | x <- [0..len - 1], y <- [0..len - 1], getCell xss (x,y) == 0
            -- work only on empty cells!
            ]

solveSudokuIntern :: [[Int]] -> [((Int,Int),IntSet)] -> [[Int]]
solveSudokuIntern xss [] = xss -- finished!
solveSudokuIntern xss possibilitiesMap =
    if List.null solutionList then [] else head solutionList
    where
        -- get Sudoku field with the least possibilities
        ((r,c),possibilitiesList) = List.minimumBy (
            \(_,a) (_,b) -> compare (IntSet.size a) (IntSet.size b)
            ) possibilitiesMap
        -- remove this field from our possibilitiesMap
        shorterPossibilitiesMap = List.delete ((r,c),possibilitiesList) possibilitiesMap
        -- generate solutions lazily (!)
        solutionList = List.take 1 $ List.filter (not . List.null) [
            let
                newPossibilitiesMap = removeElementFromSurrounding xss (r,c)
                    shorterPossibilitiesMap selectedElement
                newXss = setCell xss (r,c) selectedElement
            in
                solveSudokuIntern newXss newPossibilitiesMap
            | selectedElement <- IntSet.toList possibilitiesList
            ]

getCell :: [[Int]] -> (Int,Int) -> Int
getCell xss (r,c) = selectRow xss r !! c

getSquareFromIndex :: [[Int]] -> (Int,Int) -> [Int]
getSquareFromIndex xss (r,c) =
    concat [
        List.take num $ List.drop col $ selectRow xss j
        | j <- [row .. row + num - 1]
        ]
    where
        num = intRoot (length xss) -- number of squares per row / column
        row = (r `div` num) * num -- get index of first row of square
        col = (c `div` num) * num -- get index of first column of square

{-
Get the indices of the elements that are in the same row, column and square
as the given index (x,y) - which is part of the resulting set, too.
-}
getSurroundingIndices :: [[Int]] -> (Int,Int) -> Set (Int,Int)
getSurroundingIndices xss (r,c) =
    Set.union
        (Set.fromAscList [ (r,a) | a <- [0..len - 1] ]) -- same row
        (Set.union
            -- same column
            (Set.fromAscList [ (b,c) | b <- [0..len - 1] ])
            -- same square
            (Set.fromAscList [ (j,k) | j <- [row..row + dim - 1], k <- [col..col + dim - 1] ]))
    where
        len = length xss
        dim = intRoot len
        row = (r `div` dim) * dim -- get index of first row of square
        col = (c `div` dim) * dim -- get index of first column of square

{-
Calculate all possible choices for one index.
-}
possibilitiesForIndex :: [[Int]] -> (Int, Int) -> IntSet
possibilitiesForIndex xss (r,c) =
    IntSet.fromAscList [1..len]
        IntSet.\\
            IntSet.union
            (IntSet.fromList $ selectRow xss r)
            (IntSet.union
                (IntSet.fromList $ selectColumn xss c)
                (IntSet.fromList $ getSquareFromIndex xss (r,c)))
    where
        len = length xss
        dimension = intRoot len

{-
Remove chosen element as possibility from all fields on the same row, on
the same column and in the same square.
-}
removeElementFromSurrounding :: [[Int]] -> (Int,Int) -> [((Int,Int),IntSet)] -> Int -> [((Int,Int),IntSet)]
removeElementFromSurrounding xss (r,c) possibilitiesMap selectedElement =
    [
        (t,if Set.member t surroundingIndices then IntSet.delete selectedElement set else set)
        | (t,set) <- possibilitiesMap
        ]
    where
        surroundingIndices = getSurroundingIndices xss (r,c)

{-TTEW-}

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 $ List.map showRow xss
    where
        size = length xss
        squareSize = intRoot size 
        numberSize = size `div` 10 + 1

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

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