module Exercise03 where

import Text.Printf (printf)
import Data.List (intercalate, nub, minimumBy, find, partition, sortOn)
import Data.Maybe (fromMaybe)
import Data.Ord (comparing)
import Data.Bifunctor(second)

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

-- HA 3.1a) ii
selectColumn :: [[Int]] -> Int -> [Int]
selectColumn xss i = map (!!i) 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 = concatMap (drop col.take(col + n)) $ drop row $ take (row + n) xss
  where
    n = intRoot $ length xss
    row = (i `div` n) * n
    col = (i `mod` n) * n

-- HA 3.1b)
isValidSubsection :: [Int] -> Bool
isValidSubsection xs = nub ys == ys
  where
    ys = filter (/=0) xs

isValidSudoku :: [[Int]] -> Bool
isValidSudoku xss = all isValidSubsection $ xss ++ map (selectColumn xss) indexes ++ map (selectSquare xss) indexes
  where
    indexes = [0..length xss - 1]

insertAt :: [a] -> Int -> a -> [a]
insertAt xs i n = take i xs ++ [n] ++ drop (i + 1) xs

-- HA 3.1c)
setCell :: [[Int]] -> (Int,Int) -> Int -> [[Int]]
setCell xss (j, k) x = insertAt xss j $ insertAt (xss !! j) k x
setCell2 :: [[Int]] -> ((Int,Int), Int) -> [[Int]]
setCell2 xss (pos, x) = setCell xss pos x

-- my implementation of a list zipper, turns out i don't even use it ;)
type ZippedList a = ([a], [a])
toZipped :: [a] -> ZippedList a
toZipped (x:xs) = (xs, [x])
toList :: ZippedList a -> [a]
toList (xs, ys) = reverse ys ++ xs
forward :: ZippedList a -> ZippedList a
forward (x:xs, ys) = (xs, x:ys)
back :: ZippedList a -> ZippedList a
back (xs, y:ys) = (y:xs, ys)
nforward :: Int -> ZippedList a -> ZippedList a
nforward n zlist = iterate forward zlist !! n
nback :: Int -> ZippedList a -> ZippedList a
nback n zlist = iterate back zlist !! n
get :: ZippedList a -> a
get (xs, y:ys) = y
ins :: a -> ZippedList a -> ZippedList a
ins y (xs, ys) = (xs, y:ys)
del :: ZippedList a -> ZippedList a
del (xs, _:ys) = (xs, ys)
repl :: a -> ZippedList a -> ZippedList a
repl y zlist = ins y $ del zlist
index :: ZippedList a -> Int
index (_, ys) = length ys - 1

-- HA 3.1d)
{-WETT-}

type Pos = (Int, Int)
type EmptyCell = (Pos, [Int])
type CellSolution = (Pos, Int)

posToSquare :: Int -> Pos -> Int
posToSquare n (j, k) = (j `div` n) * n + (k `div` n)

-- checks all solutions, that are left for a cell at given position
getCellSolutions :: [[Int]] -> Pos -> [Int]
getCellSolutions xss (j, k) = filter (\val -> val `notElem` selectRow xss j && val `notElem` selectColumn xss k && val `notElem` selectSquare xss s) [1..length xss]
  where
    l = length xss
    n = intRoot l
    s = posToSquare n (j, k)

-- updates the possible solutions for each emptyCell it's given
updateCells :: Int -> [EmptyCell] -> CellSolution -> [EmptyCell]
updateCells n emptyCells ((row, col), val) = result
  where
    sqr = posToSquare n (row, col)
    isAffected emptyCell = fst (fst emptyCell) == row || snd (fst emptyCell) == col || posToSquare n (fst emptyCell) == sqr
    parts = partition isAffected emptyCells
    result = snd parts ++ map (second (filter (/= val))) (fst parts)

generateSolution :: [[Int]] -> [EmptyCell] -> [CellSolution] -> [CellSolution]
generateSolution xss emptyCells solutions
  | null emptyCells     = solutions
  | null possibleValues = [] 
  | otherwise           = fromMaybe [] maybeResult
  where
    n = intRoot $ length xss
    (pos, possibleValues) = (minimumBy $ comparing $ length.snd) emptyCells
    newEmptyCells = filter ((/= pos).fst) emptyCells
    newVals = [(updateCells n newEmptyCells (pos, val), (pos, val):solutions) | val <- possibleValues]
    maybeResult = find (/= []) $ map (uncurry $ generateSolution xss) newVals

solveSudoku :: [[Int]] -> [[Int]]
solveSudoku xss
  | not $ isValidSudoku xss                 = []
  | null solution && not (null emptyCells)  = []
  | otherwise                               = result
    where
      rng = [0..length xss-1]
      emptyCells = sortOn (length . snd) $ [((row, col), getCellSolutions xss (row, col)) | row <- rng, col <- rng, xss!!row!!col == 0]
      solution = generateSolution xss emptyCells []
      result = foldl setCell2 xss solution

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

chefsudoku :: [[Int]]
chefsudoku =
  [[0,0,0, 0,0,5, 0,8,0],
   [0,0,0, 6,0,1, 0,4,3],
   [0,0,0, 0,0,0, 0,0,0],

   [0,1,0, 5,0,0, 0,0,0],
   [0,0,0, 1,0,6, 0,0,0],
   [3,0,0, 0,0,0, 0,0,5],

   [5,3,0, 0,0,0, 0,6,1],
   [0,0,0, 0,0,0, 0,0,4],
   [0,0,0, 0,0,0, 0,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)