module Exercise03 where

import Text.Printf (printf)
import Data.List (intercalate,intersect,nub,and,minimumBy)
import qualified Data.Map as Map
import Data.Set (Set)
import qualified Data.Set as Set

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

-- 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 = [z | (ns,x) <- nss, x >= sr, x < sr + sq, (z,y) <- ns, y >= sc, y < sc + sq]  
  where 
    n = length xss
    sq = intRoot n
    sr = div i sq * sq 
    sc = mod i sq * sq
    nss = zip (map (`zip` [0..n-1]) xss) [0..n-1] 


-- HA 3.1b)
isValidSubsection :: [Int] -> Bool
isValidSubsection xs = length (nub xwzs) == length xwzs
  where 
    n = length xs
    xwzs = intersect xs [1..n]

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


-- HA 3.1c)
setCell :: [[Int]] -> (Int,Int) -> Int -> [[Int]]
setCell xss (j, k) x = take j xss ++ [ns] ++ drop (j+1) xss 
  where 
    xs = selectRow xss j
    ns = take k xs ++ [x] ++ drop (k+1) xs 

-- HA 3.1d)
{-WETT-}

sqToMap :: [[Int]] -> Map.Map (Int, Int) Int
sqToMap xss = Map.unions $ map Map.fromList nss
  where 
    n = length xss
    n' = n-1
    nss' = zip [0..n'] xss
    nss = [zip [(i,k) | k <- [0..n']] xs | (i,xs) <- nss']

sR :: Int -> Int -> Map.Map (Int, Int) Int -> [Int]
sR sz idx mx = [mx Map.! (idx,k) | k <- [0..sz]] 

sC :: Int -> Int -> Map.Map (Int, Int) Int -> [Int]
sC sz idx mx = [mx Map.! (k,idx) | k <- [0..sz]]

sS :: Int -> Int -> Map.Map (Int, Int) Int -> [Int]
sS sz idx mx = [mx Map.! (k,j) | k <- [sx..ex], j <- [sy..ey]]
  where
    sq = intRoot $ sz + 1
    sx = div idx sq * sq 
    sy = mod idx sq * sq 
    ex = sx + sq - 1
    ey = sy + sq - 1

cS :: (Int, Int) -> Int -> Map.Map (Int, Int) Int -> Map.Map (Int, Int) Int
cS xx v mx = Map.insert xx v mx

notIn :: Int -> Set.Set Int -> Set.Set Int
notIn sz ss = ns Set.\\ ss 
  where 
    ns = Set.fromAscList [1..sz+1]

toS :: Int -> (Int, Int) -> Int
toS sz (x,y) = sn
  where
    sq = intRoot $ sz + 1
    sx = div x sq * sq
    sy = div y sq 
    sn = sx + sy 

fPos :: Int -> Map.Map (Int, Int) Int -> (Int, Int) -> Set.Set Int 
fPos sz mx (x,y) = notIn sz $ Set.unions $ map Set.fromList (sR sz x mx : sC sz y mx : sS sz sn mx : [])  
  where
    sn = toS sz (x,y) 


isV :: Int -> Map.Map (Int, Int) Int -> Bool
isV sz mx = and [(isValidSubsection $ sR sz i mx) && (isValidSubsection $ sC sz i mx)  &&  (isValidSubsection $ sS sz i mx) | i <- [0..sz]]


findZ :: [[Int]] -> [(Int,Int)]
findZ xss = [(x,y) | (xs,x) <- nss, (z,y) <- xs, z == 0]
  where
    n = length xss
    nss = zip (map (`zip` [0..n-1]) xss) [0..n-1]

dZ :: Int -> [(Int,Int)] -> Map.Map (Int, Int) Int -> Map.Map (Int, Int) (Set.Set Int)
dZ sz zs mx = Map.fromList (zip zs (map (fPos sz mx) zs))

mapToSq :: Map.Map (Int, Int) Int -> [[Int]]
mapToSq mx = [sR n i mx | i <- [0..n]]
  where
    n = fst $ fst $ Map.findMax mx

solveSudoku :: [[Int]] -> [[Int]]
solveSudoku xss 
  | not (isV n mx) = []
  | pos = mapToSq erg
  | otherwise = []
  where
    n = length xss - 1
    mx = sqToMap xss
    (erg,pos) = solveSudokuM n mx (dZ n (findZ xss) mx)


comp :: ((Int, Int),(Set.Set Int)) -> ((Int, Int),(Set.Set Int)) -> Ordering
comp (_,sx) (_,sy) = compare (Set.size sx) (Set.size sy) 
 
findBest :: Map.Map (Int, Int) (Set.Set Int) -> (Int, Int)
findBest msx = fst (minimumBy comp $ Map.toList msx)

genR :: Int -> (Int,Int) -> [(Int,Int)]
genR sz (x,y) = [(x,i) | i <- [0..sz]]

genC :: Int -> (Int,Int) -> [(Int,Int)]
genC sz (x,y) = [(i,y) | i <- [0..sz]]

genS :: Int -> (Int,Int) -> [(Int,Int)]
genS sz (x,y) = [(i,j) | i <- [sx..ex], j <- [sy..ey]]
  where
    ss = toS sz (x,y)
    sq = intRoot (sz+1)
    sx = div ss sq * sq
    ex = sx + sq - 1
    sy = mod ss sq * sq
    ey = sy + sq - 1
    

alterS :: [(Int,Int)] -> Int -> Map.Map (Int,Int) (Set.Set Int) -> Map.Map (Int,Int) (Set.Set Int)
alterS [] _ msx = msx 
alterS ((x,y):xs) ch msx = alterS xs ch (Map.insert (x,y) ns msx) 
  where
    os = msx Map.! (x,y)
    ns = Set.delete ch os

fCh :: Int -> Map.Map (Int, Int) Int -> (Int,Int) -> [(Int,Int)]
fCh sz mx (x,y) = [(x',y')| (x',y') <- al,mx Map.! (x',y') == 0]
  where
    al = genR sz (x,y) ++ (genC sz (x,y)) ++ (genS sz (x,y))

solveSudokuMB :: Int -> Map.Map (Int, Int) Int -> Map.Map (Int,Int) (Set.Set Int) -> (Int,Int) -> [Int] -> (Map.Map (Int, Int) Int, Bool)
solveSudokuMB _ mx _ _ [] = (mx,False)
solveSudokuMB sz mx msx (x,y) (z:zs)
  | pos = (erg,pos)
  | otherwise = solveSudokuMB sz mx msx (x,y) zs
  where 
    nmx = Map.insert (x,y) z mx
    tt = fCh sz nmx (x,y)
    nmsx = Map.delete (x,y) (alterS tt z msx)
    (erg,pos) = solveSudokuM sz nmx nmsx



solveSudokuM :: Int -> Map.Map (Int, Int) Int -> Map.Map (Int,Int) (Set.Set Int) -> (Map.Map (Int, Int) Int, Bool)
solveSudokuM sz mx msx
  | Map.size msx == 0 = (mx,True)
  | otherwise = solveSudokuMB sz mx msx x pp
  where 
   x = findBest msx
   pp = Set.toList (msx Map.! x)


{-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 $ 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)