module Exercise03 where

import Text.Printf (printf)
import Data.List (intercalate, sort, nub)
import Data.Map as Map
import Data.Set as Set
    ( Set,
      (\\),
      empty,
      fromAscList,
      notMember,
      singleton,
      size,
      toList,
      union )
import Data.List as List

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

-- HA 3.1a) ii
selectColumn :: [[Int]] -> Int -> [Int]
selectColumn xss i = [col !! i | col <- 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 = 
  let l = intRoot (length xss)
      x = (i `mod` l) * l
      y = (i `div` l) * l
  in concat [Prelude.take l $ Prelude.drop x $ selectRow xss (y + index) | index <- [0..(l - 1)]]

-- HA 3.1b)
isValidSubsection :: [Int] -> Bool
isValidSubsection xss =
  let filtered = Prelude.filter (/= 0) xss
  in length (nub filtered) == length filtered

isValidSudoku :: [[Int]] -> Bool
isValidSudoku xss = and [
  isValidSubsection (selectRow xss index) && 
  isValidSubsection (selectColumn xss index) &&
  isValidSubsection (selectSquare xss index) | index <- [0..(length xss - 1)]]

-- HA 3.1c)
setCell :: [[Int]] -> (Int,Int) -> Int -> [[Int]]
setCell xss (y, x) n =
  let upper = Prelude.take y xss
      lower = Prelude.drop (y + 1) xss
      row = selectRow xss y
      modified_row = Prelude.take x row ++ [n] ++ Prelude.drop (x + 1) row
  in upper ++ [modified_row] ++ lower

{-WETT-}

-- format idea: https://abhinavsarkar.net/posts/fast-sudoku-solver-in-haskell-1/
-- (without copying the implementation ofc)

-- HA 3.1d)
solveSudoku :: [[Int]] -> [[Int]]
solveSudoku xss = mapToSudoku $ backtrack $ sudokuToMap xss

-- custom functions

type Sudoku = Map Int (Set Int)

sudokuToMap :: [[Int]] -> Sudoku
sudokuToMap sudoku =
  let dim = length sudoku
      conc = [0..] `zip` concat sudoku 
  in Prelude.foldl (\acc (index, value) -> Map.insert index (if value == 0 then Set.fromAscList [1..dim] else Set.singleton value) acc) Map.empty conc

mapToSudoku :: Sudoku -> [[Int]]
mapToSudoku sudoku = [Prelude.map (\i -> head $ Set.toList (sudoku ! i)) r | r <- rows sudoku]
--

rows :: Sudoku -> [[Int]]
rows sudoku = [Prelude.map (\x -> x + dim * cell) [0..(dim - 1)] | cell <- [0..(dim - 1)]]
  where dim = intRoot $ Map.size sudoku

cells :: Sudoku -> [[Int]]
cells sudoku = [Prelude.map (\x -> x * dim + row) [0..(dim - 1)] | row <- [0..(dim - 1)]]
  where dim = intRoot $ Map.size sudoku

squares :: Sudoku -> [[Int]]
squares sudoku = [concat [Prelude.map (\x -> (vo * sqSize + vi) * dim + ho * sqSize + x) sqList| vi <- sqList] | ho <- sqList, vo <- sqList]
  where 
    sqSize = intRoot dim
    sqList = [0..(sqSize - 1)]
    dim = intRoot $ Map.size sudoku

cleanPossibilites :: Sudoku -> Sudoku
cleanPossibilites sudoku
  | duplSingle sudoku = Map.empty
  | otherwise =
      let cleaned = cleanPHelper sudoku (rows sudoku ++ cells sudoku ++ squares sudoku)
      in
        if cleaned == sudoku then sudoku
        else cleanPossibilites cleaned

cleanPHelper :: Sudoku -> [[Int]] -> Sudoku
cleanPHelper = Prelude.foldl cleanArea

cleanArea :: Sudoku -> [Int] -> Sudoku
cleanArea sudoku is =
  let singles = getSingles sudoku is
  in Prelude.foldl (\s i -> if Set.size (sudoku ! i) == 1 then s else Map.insert i ((sudoku ! i) Set.\\ singles) s) sudoku is

getSingles :: Sudoku -> [Int] -> Set Int
getSingles sudoku is =
  Prelude.foldl (\acc set -> if Set.size set == 1 then Set.union set acc else acc) Set.empty $ Prelude.map (sudoku !) is

-- actual backtracking

backtrack :: Sudoku -> Sudoku
backtrack sudoku = backtrackHelp sudoku 0

backtrackHelp :: Sudoku -> Int -> Sudoku
backtrackHelp sudoku i
  | Map.null sudoku = Map.empty
  | solved sudoku = sudoku
  | i >= Map.size sudoku = Map.empty
  | otherwise = backtrackHelp2 sudoku i (Set.toList $ sudoku ! i)

backtrackHelp2 :: Sudoku -> Int -> [Int] -> Sudoku
backtrackHelp2 _ _ [] = Map.empty
backtrackHelp2 sudoku i (x:xs) =
  let modifiedSudoku = Map.insert i (Set.singleton x) sudoku
      cleaned = cleanPossibilites modifiedSudoku
      try = backtrackHelp cleaned (i + 1)
  in
    if Map.null try then backtrackHelp2 sudoku i xs
    else try

duplSingle :: Sudoku -> Bool
duplSingle sudoku =
  let r = rows sudoku
      c = cells sudoku
      sqr = squares sudoku
  in not $ Prelude.null [0 | is <- r ++ c ++ sqr, dupl sudoku is]

dupl :: Sudoku -> [Int] -> Bool
dupl sudoku is =
  let list = Prelude.filter (/= 0) $ Prelude.map (head . Set.toList) $ Prelude.filter (\set -> Set.size set == 1) $ Prelude.map (sudoku !) is
  in length (List.nub list) /= length list

filled :: Sudoku -> Bool
filled sudoku = all
    (\row -> and [Set.size (sudoku ! i) == 1 && 0 `Set.notMember` (sudoku ! i) | i <- row])
    (rows sudoku)

solved :: Sudoku -> Bool
solved sudoku
  | not $ filled sudoku = False
  | not $ solvedPart sudoku $ rows sudoku = False
  | not $ solvedPart sudoku $ cells sudoku = False
  | not $ solvedPart sudoku $ squares sudoku = False
  | otherwise = True

solvedPart :: Sudoku -> [[Int]] -> Bool
solvedPart sudoku xss = and [Set.fromAscList [1..(intRoot $ Map.size sudoku)] == Prelude.foldl (\ acc x -> acc `Set.union` (sudoku ! x)) Set.empty r | r <- xss]

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

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

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