-- vim: set tabstop=2 softtabstop=0 expandtab shiftwidth=2:

{-# LANGUAGE BangPatterns #-}

module Exercise03 where
--module Main where

import Text.Printf (printf)
import Data.List
import Data.Functor (fmap)
import Data.Bits
import Data.Tuple
import Data.Ord
import Data.Foldable
import Data.Maybe
import Data.Containers.ListUtils
import qualified Data.Sequence as Seq

-- 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 (take q . drop (q * d2)) $ take q $ drop (q * d1) xss
  where
    q = intRoot $ length xss
    (d1,d2) = i `divMod` q

-- HA 3.1b)
--isValidSubsectionSlow :: [Int] -> Bool
--isValidSubsectionSlow xs = length entries == length (nubOrd entries)
--  where entries = filter (/=0) xs

-- Fast version
isValidSubsection :: [Int] -> Bool
isValidSubsection xs = let (ss,nzc) = f xs in popCount (shiftR ss 1) == nzc
  where
    f :: [Int] -> (Int, Int)
    f [] = (0, 0)
    f (0:rs) = f rs
    f (x:rs) = let (a,b) = f rs in (a `setBit` x, succ b)

-- Check if a sudoku is completely valid, with cool <*> operator
isValidSudoku :: [[Int]] -> Bool
isValidSudoku xss = all isValidSubsection $ [selectRow xss, selectColumn xss, selectSquare xss] <*> [0..pred $ length xss]


-- Only check affected rows
isStillValidSudoku :: [[Int]] -> (Int,Int) -> Bool
isStillValidSudoku xss (x,y) =  isValidSubsection (selectRow xss y)
                             && isValidSubsection (selectColumn xss x)
                             && isValidSubsection (selectSquare xss (3 * (y `div` 3) + (x `div` 3)))

-- HA 3.1c)
setCell :: [[Int]] -> (Int,Int) -> Int -> [[Int]]
setCell xss (k, j) x
    = toList
    $ fmap toList
    $ Seq.chunksOf n
    $ Seq.update (k * n + j) x
    $ Seq.fromList
    $ concat xss
  where n = length xss

-- HA 3.1d)

{-WETT-}

-- Seqdoku = Flattened sudoku with Data.Sequence.Seq as list type
type Seqdoku = Seq.Seq Int
-- Idoku = Seqdoku plus integer denoting the sudoku side-length
type Idoku = (Seqdoku, Int)
-- Constraindoku = Constraints for each cell of seqdoku
type Constraindoku = Seq.Seq Int

ctz :: Int -> Int
ctz = countTrailingZeros

-- Convert original double list sudoku to idoku
badSudokuToIdoku :: [[Int]] -> Idoku
badSudokuToIdoku xss = (Seq.fromList $ concat xss, length xss)

-- Convert Maybe idoku to original double list sudoku, or empty list
idokuToBadSudoku :: Maybe Idoku -> [[Int]]
idokuToBadSudoku Nothing = []
idokuToBadSudoku (Just (seqdoku,n)) = toList $ toList <$> Seq.chunksOf n seqdoku

-- Fast boi to set a cell
setIdokuCell :: Idoku -> Int -> Int -> Idoku
setIdokuCell (seqdoku,n) i val = (Seq.update i val seqdoku, n)

-- Check if index is relevant to a reference cell
relevantIndices :: Int -> Int -> [Int]
relevantIndices !n !i = seqdokuSquareIndices n squid ++ ([\k -> r*n+k, \k -> n*k+c] <*> [0..pred n])
  where
    (r,c) = i `divMod` n
    q = intRoot n
    squid = q * (r `div` q) + (c `div` q)

-- Create list with selected row values
selectIdokuRow :: Idoku -> Int -> [Int]
selectIdokuRow (!seqdoku,!n) i = [Seq.index seqdoku (i*n+k) | !k <- [0..pred n]]

-- Create list with selected column values
selectIdokuColumn :: Idoku -> Int -> [Int]
selectIdokuColumn (!seqdoku,!n) i = [Seq.index seqdoku (n*k+i) | !k <- [0..pred n]]

-- Create list with selected square values, by looking up the
-- correct index from an autogenereated index list
selectIdokuSquare :: Idoku -> Int -> [Int]
selectIdokuSquare (!seqdoku,!n) i = [Seq.index seqdoku k | !k <- seqdokuSquareIndices n i]

-- Given side length n and square index i this returns the flattened indices for all numbers
-- in the square
seqdokuSquareIndices :: Int -> Int -> [Int]
seqdokuSquareIndices 1 0 = [0]
seqdokuSquareIndices 4 0 = [0,1,4,5]
seqdokuSquareIndices 4 1 = [2,3,6,7]
seqdokuSquareIndices 4 2 = [8,9,12,13]
seqdokuSquareIndices 4 3 = [10,11,14,15]
seqdokuSquareIndices 9 0 = [0,1,2,9,10,11,18,19,20]
seqdokuSquareIndices 9 1 = [3,4,5,12,13,14,21,22,23]
seqdokuSquareIndices 9 2 = [6,7,8,15,16,17,24,25,26]
seqdokuSquareIndices 9 3 = [27,28,29,36,37,38,45,46,47]
seqdokuSquareIndices 9 4 = [30,31,32,39,40,41,48,49,50]
seqdokuSquareIndices 9 5 = [33,34,35,42,43,44,51,52,53]
seqdokuSquareIndices 9 6 = [54,55,56,63,64,65,72,73,74]
seqdokuSquareIndices 9 7 = [57,58,59,66,67,68,75,76,77]
seqdokuSquareIndices 9 8 = [60,61,62,69,70,71,78,79,80]
seqdokuSquareIndices 16 0 = [0,1,2,3,16,17,18,19,32,33,34,35,48,49,50,51]
seqdokuSquareIndices 16 1 = [4,5,6,7,20,21,22,23,36,37,38,39,52,53,54,55]
seqdokuSquareIndices 16 2 = [8,9,10,11,24,25,26,27,40,41,42,43,56,57,58,59]
seqdokuSquareIndices 16 3 = [12,13,14,15,28,29,30,31,44,45,46,47,60,61,62,63]
seqdokuSquareIndices 16 4 = [64,65,66,67,80,81,82,83,96,97,98,99,112,113,114,115]
seqdokuSquareIndices 16 5 = [68,69,70,71,84,85,86,87,100,101,102,103,116,117,118,119]
seqdokuSquareIndices 16 6 = [72,73,74,75,88,89,90,91,104,105,106,107,120,121,122,123]
seqdokuSquareIndices 16 7 = [76,77,78,79,92,93,94,95,108,109,110,111,124,125,126,127]
seqdokuSquareIndices 16 8 = [128,129,130,131,144,145,146,147,160,161,162,163,176,177,178,179]
seqdokuSquareIndices 16 9 = [132,133,134,135,148,149,150,151,164,165,166,167,180,181,182,183]
seqdokuSquareIndices 16 10 = [136,137,138,139,152,153,154,155,168,169,170,171,184,185,186,187]
seqdokuSquareIndices 16 11 = [140,141,142,143,156,157,158,159,172,173,174,175,188,189,190,191]
seqdokuSquareIndices 16 12 = [192,193,194,195,208,209,210,211,224,225,226,227,240,241,242,243]
seqdokuSquareIndices 16 13 = [196,197,198,199,212,213,214,215,228,229,230,231,244,245,246,247]
seqdokuSquareIndices 16 14 = [200,201,202,203,216,217,218,219,232,233,234,235,248,249,250,251]
seqdokuSquareIndices 16 15 = [204,205,206,207,220,221,222,223,236,237,238,239,252,253,254,255]
seqdokuSquareIndices !n !squareIndex = concatMap (\x -> [fn x..fn x+q']) [0..q']
  where
    !q = intRoot n
    !q' = pred q
    (!dy,!dx) = squareIndex `divMod` q
    fn k = n * (k+q*dy) + q*dx

-- Fast version. We use the cpu instructions popCount and countTrailingZeros for ultimate speed.
-- HINT: The MC told me that the algorithm only has to work for sudokus with a size of equal or less
--       than (7*7)x(7*7). This allows me to use Int as bitfields instead of more expensive Integer types.
isValidIdokuSubsection :: [Int] -> Bool
isValidIdokuSubsection !xs = let (ss,nzc) = f xs in popCount (shiftR ss 1) == nzc
  where
    f :: [Int] -> (Int, Int)
    f [] = (0, 0)
    f (0:rs) = f rs
    f (x:rs) = let (a,b) = f rs in (a `setBit` x, succ b)

-- Check if a idoku is completely valid
isValidIdoku :: Idoku -> Bool
isValidIdoku idoku@(_,n) = all isValidIdokuSubsection
                             $ [selectIdokuRow, selectIdokuColumn, selectIdokuSquare]
                               <*> [idoku]
                               <*> [0..pred n]

-- Check if a idoku is still valid after changing the value at the given position
isStillValidIdoku :: Idoku -> (Int,Int) -> Bool
isStillValidIdoku idoku@(_,n) (x,y)
  =  isValidIdokuSubsection (selectIdokuRow idoku y)
  && isValidIdokuSubsection (selectIdokuColumn idoku x)
  && isValidIdokuSubsection (selectIdokuSquare idoku (q * (y `div` q) + (x `div` q)))
  where q = intRoot n

-- Returns flat-index of first 0
findNextFreeCell :: Idoku -> Maybe Int
findNextFreeCell (seqdoku,n) = Seq.findIndexL (==0) seqdoku

-- Returns flat-index of the 0 with most constraints
findNextFreeCellMostConstraints :: Idoku -> Constraindoku -> Maybe Int
findNextFreeCellMostConstraints (seqdoku,n) constraints =
  case indices of
    [] -> Nothing
    xs -> Just $ minimumBy (comparing (\k -> popCount (constraints `Seq.index` k))) xs
  where
    indices = Seq.findIndicesL (==0) seqdoku

-- Solve idoku with backtracking
solveIdokuBacktracking :: Idoku -> Maybe Idoku
solveIdokuBacktracking idoku@(_,n) =
  case possidokus of
    -- If there are no possible new idokus, we cannot solve it and return Noting
    [] -> Nothing
    -- Otherwise, we return the first solved subsudoku
    (x:_) -> Just x
  where
    possidokus = case findNextFreeCell idoku of
      -- No free cell. If the current idoku is valid, we have found a solution.
      -- Otherwise, there is none for this path.
      Nothing -> [idoku | isValidIdoku idoku]
      -- If there is a free cell, return the list of solved subsudokus,
      -- removing the dead ends.
      Just freeCell -> catMaybes [solveIdokuBacktracking newdoku | num <- [1..n],
                                    -- Create new idoku for current path
                                    newdoku <- [setIdokuCell idoku freeCell num],
                                    -- Prefilter to remove obviously invalid ones
                                    isStillValidIdoku newdoku $ swap (freeCell `divMod` n)]

removePossiblity :: Int -> Int -> Int
removePossiblity constraint i = constraint `clearBit` i

addPossiblity :: Int -> Int -> Int
addPossiblity constraint i = constraint `setBit` i

anyConstraint :: Int -> Int
anyConstraint n = 2^succ n - 2 -- "-2" instead of "-1" also clears the 0th bit, which is never an option.

-- Checks for empty (unsolvable) constraints
isConstraintUnsolvable :: Constraindoku -> Bool
isConstraintUnsolvable constraindoku =
  case Seq.elemIndexL 0 constraindoku of
    Nothing -> False
    Just a -> True

calcConstraint :: Idoku -> Int -> Int
calcConstraint idoku@(seqdoku,n) i =
  case cell of
    0 -> foldl removePossiblity (anyConstraint n) (row ++ col ++ square)
    _ -> addPossiblity 0 cell
  where
    cell   = seqdoku `Seq.index` i
    relevant = relevantIndices n i
    q      = intRoot n
    (y,x)  = i `divMod` n
    row    = selectIdokuRow idoku y
    col    = selectIdokuColumn idoku x
    square = selectIdokuSquare idoku (q * (y `div` q) + (x `div` q))

-- SLOW! Calculate initial constraints for sudoku
--initConstraints :: Idoku -> Constraindoku
--initConstraints idoku@(seqdoku,n) =
--  Seq.fromList $ [calcConstraint idoku] <*> [0..pred $ n*n]

-- FAST! Calculate initial constraints for sudoku
initConstraints :: Idoku -> Constraindoku
initConstraints idoku@(seqdoku,n) =
  Seq.fromList $ [consCell] <*> [0..pred $ n*n]
  where
    !consRows = [foldl removePossiblity (anyConstraint n) (selectIdokuRow idoku k) | k <- [0..pred n]]
    !consCols = [foldl removePossiblity (anyConstraint n) (selectIdokuColumn idoku k) | k <- [0..pred n]]
    !consSqus = [foldl removePossiblity (anyConstraint n) (selectIdokuSquare idoku k) | k <- [0..pred n]]
    consCell i = case seqdoku `Seq.index` i of
        0 -> (consRows !! y) .&. (consCols !! x) .&. (consSqus !! s)
        num -> addPossiblity 0 num
      where
        q      = intRoot n
        (y,x)  = i `divMod` n
        s      = (q * (y `div` q) + (x `div` q))

-- Transfer constraints with only one possibility to the idoku
transferSingleConstraints :: Idoku -> Constraindoku -> Idoku
transferSingleConstraints (seqdoku,n) constraindoku =
  (fmap process (Seq.zip seqdoku constraindoku), n)
  where
    process :: (Int, Int) -> Int
    process (cell,constraint) = if popCount constraint == 1 then ctz constraint
                                                            else cell

countEmpty :: Idoku -> Int
countEmpty (seqdoku,_) = foldl (\x y -> x + if y == 0 then 1 else 0) 0 seqdoku

possibleNumbersFor :: Constraindoku -> Int -> [Int]
possibleNumbersFor constraints cell = aux c 0
  where
    c = Seq.index constraints cell
    aux 0 _ = []
    aux n k = (k + trailingZeros) : aux (shiftR n processedBits) (k + processedBits)
      where
        trailingZeros = ctz n
        processedBits = trailingZeros + 1

calculateNewConstraints :: Int -> Constraindoku -> Int -> Int -> Constraindoku
calculateNewConstraints n constraints fixedCell selectedNumber =
  Seq.update fixedCell (addPossiblity 0 selectedNumber)
    (Seq.mapWithIndex updateRelevantConstraints constraints)
  where
    relevant = relevantIndices n fixedCell
    updateRelevantConstraints idx co
      | idx `elem` relevant = removePossiblity co selectedNumber
      | otherwise = co

-- Solve idoku with backtracking
solveIdokuConstraintTrackingBacktrack :: Idoku -> Constraindoku -> Maybe Idoku
solveIdokuConstraintTrackingBacktrack idoku@(_,n) constraints =
  case possidokus of
    -- If there are no possible new idokus, we cannot solve it and return Noting
    [] -> Nothing
    -- Otherwise, we return the first solved subsudoku
    (x:_) -> Just x
  where
    possidokus = case findNextFreeCell idoku of
      -- No free cell. If the current idoku is valid, we have found a solution.
      -- Otherwise, there is none for this path.
      Nothing -> [idoku | isValidIdoku idoku]
      -- If there is a free cell, return the list of solved subsudokus,
      -- removing the dead ends.
      Just freeCell -> catMaybes [solveConstrainedIdoku newdoku (calculateNewConstraints n constraints freeCell num) | num <- possibleNumbersFor constraints freeCell,
                                    -- Create new idoku for current path
                                    newdoku <- [setIdokuCell idoku freeCell num],
                                    -- Prefilter to remove obviously invalid ones
                                    isStillValidIdoku newdoku $ swap (freeCell `divMod` n)]

-- Solve idoku with basic constraint tracking
solveConstrainedIdoku :: Idoku -> Constraindoku -> Maybe Idoku
solveConstrainedIdoku idoku@(seqdoku,n) constraints =
      if isConstraintUnsolvable constraints then
        Nothing
      else let transferdoku = transferSingleConstraints idoku constraints in
        if countEmpty transferdoku <= 16 then
          solveIdokuBacktracking transferdoku
        else
          let new_constraints = initConstraints transferdoku in
          if constraints == new_constraints then
            -- We need to backtrack one cell to continue
            solveIdokuConstraintTrackingBacktrack transferdoku constraints
          else
            solveConstrainedIdoku transferdoku new_constraints

-- Solve idoku with basic constraint tracking
solveIdokuWithConstraints :: Idoku -> Maybe Idoku
solveIdokuWithConstraints idoku =
  solveConstrainedIdoku idoku (initConstraints idoku)

-- Wrap the sudoku in an idoku, solve it (maybe), and create the required representation again.
solveSudoku :: [[Int]] -> [[Int]]
--solveSudoku = idokuToBadSudoku . solveIdokuBacktracking . badSudokuToIdoku
solveSudoku = idokuToBadSudoku . solveIdokuWithConstraints . badSudokuToIdoku

{-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 [] = "Unsolvable\n"
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)