{-# LANGUAGE RankNTypes #-}
module Main where

import Control.Monad
import Control.Monad.Writer
import Data.Char
import Data.Either
import Data.List
import System.Exit

import qualified FSu


lupdate :: Int -> a -> [a] -> [a]
lupdate 0 v (_:xs) = v:xs
lupdate i v (x:xs) = x : lupdate (i-1) v xs
lupdate _ _ [] = error "Empty list in lupdate"

b2i :: Bool -> Int
b2i = fromEnum

subsets :: [a] -> [[a]]
subsets [] = [[]]
subsets (x:xs) = let s = subsets xs in s ++ map (x :) s


newtype Sudoku = Sudoku [Either Value Poss]  -- length 81
  deriving (Show, Eq)

type Value = Int  -- 0..8

newtype Poss = Poss [Bool]  -- length 9
  deriving (Show, Eq)


at :: Sudoku -> Int -> Either Value Poss
at (Sudoku l) i = l !! i

atv :: Sudoku -> Int -> Maybe Value
atv su i = case at su i of
    Left v -> Just v
    Right _ -> Nothing

atp :: Sudoku -> Int -> Poss
atp su i = case at su i of
    Left _ -> Poss (replicate 9 False)
    Right p -> p

update :: Int -> Either Value Poss -> Sudoku -> Sudoku
update i v (Sudoku l) = Sudoku $ lupdate i v l

listPoss :: [Value] -> Poss
listPoss l = Poss [i `elem` l | i <- [0..8]]

possList :: Poss -> [Value]
possList (Poss m) = map snd $ filter fst $ zip m [0..8]

possAnd :: Poss -> Poss -> Poss
possAnd (Poss l1) (Poss l2) = Poss (zipWith (&&) l1 l2)

possOr :: Poss -> Poss -> Poss
possOr (Poss l1) (Poss l2) = Poss (zipWith (||) l1 l2)

possElem :: Value -> Poss -> Bool
possElem n (Poss l) = 0 <= n && n < 9 && l !! n

possExact1 :: [Poss] -> Poss
possExact1 [] = Poss (replicate 9 False)
possExact1 ps =
    Poss $ map (== 1) $
        foldl1 (zipWith (+)) $ map (\(Poss p) -> map b2i p) ps

fullPoss :: Poss
fullPoss = Poss (replicate 9 True)

emptyPoss :: Poss
emptyPoss = Poss (replicate 9 False)

possMakeAllExcept :: Value -> Poss
possMakeAllExcept n =
    Poss $ map (\i -> if i == n then False else True) [0..8]

possGetIfInOne :: [Poss] -> Value -> Maybe Int
possGetIfInOne ps n =
    case concatMap (\(p,i) -> if n `possElem` p then [i] else [])
                   (zip ps [0..]) of
        [i] -> Just i
        _ -> Nothing


backtrackSolve :: Sudoku -> [Sudoku]
backtrackSolve (Sudoku l) = do
    let ml = map (either Just (const Nothing)) l
    res <- FSu.solve ml
    let l' = map (maybe (Right (listPoss [0..8])) Left) res
    return $ Sudoku l'


data Region = Row Int Sudoku
            | Col Int Sudoku
            | Block Int Sudoku
            | Cell Int Sudoku
  deriving (Show, Eq)

rowOf :: Region -> Region
rowOf (Cell i su) = Row (i `div` 9) su

colOf :: Region -> Region
colOf (Cell i su) = Col (i `mod` 9) su

blockOf :: Region -> Region
blockOf (Cell i su) = Block (3 * (i `div` 27) + (i `mod` 9) `div` 3) su

valueOf :: Region -> Maybe Value
valueOf (Cell i su) = su `atv` i

possOf :: Region -> Poss
possOf (Cell i su) = su `atp` i

contentsOf :: Region -> Either Value Poss
contentsOf (Cell i su) = su `at` i

indexOf :: Region -> Int
indexOf (Cell i _) = i


sudokuOf :: Region -> Sudoku
sudokuOf (Cell _ su) = su
sudokuOf (Row _ su) = su
sudokuOf (Col _ su) = su
sudokuOf (Block _ su) = su

indexRange :: Region -> [Int]
indexRange (Cell _ _) = [0]
indexRange (Row _ _) = [0..8]
indexRange (Col _ _) = [0..8]
indexRange (Block _ _) = [0..8]

-- Doesn't check sudoku equality
cellIsIn :: Region -> Region -> Bool
cellIsIn (Cell i _) (Cell j _) = i == j
cellIsIn (Cell i _) (Row r _) = i `div` 9 == r
cellIsIn (Cell i _) (Col c _) = i `mod` 9 == c
cellIsIn (Cell i _) (Block b _) =
    (i `mod` 9) `div` 3 == b `mod` 3 && i `div` 27 == b `div` 3

adaptRegion :: Region -> Sudoku -> Region
adaptRegion (Cell i _) su = Cell i su
adaptRegion (Row i _) su = Row i su
adaptRegion (Col i _) su = Col i su
adaptRegion (Block i _) su = Block i su


class HasCells a where
    cell :: Int -> a -> Region

    cells :: a -> [Region]

instance HasCells Sudoku where
    cell i su | 0 <= i && i < 81 = Cell i su

    cells su = [Cell i su | i <- [0..80]]

instance HasCells Region where
    cell i (Row r su) | 0 <= i && i < 9 = Cell (9 * r + i) su
    cell i (Col c su) | 0 <= i && i < 9 = Cell (9 * i + c) su
    cell i (Block b su) | 0 <= i && i < 9 =
        let bx = 3 * (b `mod` 3)
            by = 3 * (b `div` 3)
        in Cell (9 * (by + i `div` 3) + bx + i `mod` 3) su
    cell 0 c@(Cell _ _) = c
    cell _ _ = undefined

    cells r = [cell i r | i <- [0..8]]


class HasValues a where
    values :: a -> [Value]

instance HasValues Sudoku where
    values (Sudoku l) = lefts l

instance HasValues Region where
    values (Cell i su) = maybe [] pure (su `atv` i)
    values reg =  -- row, col, block
        concatMap values [cell i reg | i <- indexRange reg]


class UpdateValue a where
    updateValue :: Int -> Value -> a -> Sudoku

instance UpdateValue Sudoku where
    updateValue i v su = update i (Left v) su

instance UpdateValue Region where
    updateValue 0 v (Cell j su) = update j (Left v) su
    updateValue _ _ (Cell _ _) = undefined
    updateValue i v reg = updateValue 0 v (cell i reg)  -- row, col, block


class MaskPoss a where
    maskPoss :: a -> Poss -> Sudoku

instance MaskPoss Region where
    maskPoss (Cell i su) p = case su `at` i of
        Left _ -> su
        Right p' -> update i (Right (p `possAnd` p')) su
    maskPoss row@(Row   r _) p =
        sudokuOf $ foldl (\row' i -> Row   r $ maskPoss (cell i row') p) row [0..8]
    maskPoss col@(Col   c _) p =
        sudokuOf $ foldl (\col' i -> Col   c $ maskPoss (cell i col') p) col [0..8]
    maskPoss blk@(Block b _) p =
        sudokuOf $ foldl (\blk' i -> Block b $ maskPoss (cell i blk') p) blk [0..8]


-- Returns cells in first region that also occur in second region
intersectReg :: Region -> Region -> [Region]
intersectReg reg1 reg2 =
    filter (\c -> c `cellIsIn` reg2) [cell i reg1 | i <- indexRange reg1]


readSudoku :: String -> Sudoku
readSudoku s = Sudoku $ flip map (words s) $ \w -> case w of
    "." -> Right (Poss $ replicate 9 True)
    [c] | '1' <= c && c <= '9' -> Left (ord c - ord '1')
    _ -> error "Invalid sudoku input"

writeSudokuGeneric :: Sudoku -> Bool -> String
writeSudokuGeneric su full =
    intercalate "\n"
        [testinter n "\n" ++
            intercalate " "
                (map (\i -> testinter i " " ++
                            printone (su `atv` i) ++
                            printposs (su `atp` i))
                     [9*n..9*n+8])
        | n <- [0..8]]
  where
    printone :: Maybe Value -> String
    printone Nothing = "."
    printone (Just n) = show (n + 1)

    printposs :: Poss -> String
    printposs (Poss m) =
        let s = map snd $ filter fst $ zip m ['1'..'9']
            padded = replicate (9 - length s) ' ' ++ s
        in if full then "(" ++ padded ++ ")" else ""

    testinter :: Int -> String -> String
    testinter i s = let n = i `mod` 9 in if n > 0 && n `mod` 3 == 0 then s else ""

writeSudoku :: Sudoku -> String
writeSudoku su = writeSudokuGeneric su False

writeSudokuFull :: Sudoku -> String
writeSudokuFull su = writeSudokuGeneric su True

writeSudokuDiff :: Sudoku -> Sudoku -> String
writeSudokuDiff su1 su2 =
    intercalate "\n" $ flip map [0..8] $ \r ->
        let line = flip concatMap [0..8] $ \c ->
                (if c > 1 && c `mod` 3 == 0 then " " else "")
                ++ (if c > 1 then " " else "")
                ++ case (su1 `at` (9 * r + c), su2 `at` (9 * r + c)) of
                    (Left v1, Left v2)
                        | v1 == v2 -> show (v2 + 1) ++ printposs pe pe
                        | otherwise -> high (show (v2 + 1)) ++ printposs pe pe
                    (Left _, Right p) -> high ("." ++ printposs p p)
                    (Right _, Left v) -> high (show (v + 1) ++ printposs pe pe)
                    (Right p1, Right p2) -> "." ++ printposs p1 p2
        in (if r > 1 && r `mod` 3 == 0 then "\n" else "") ++ line
  where
    high :: String -> String
    high s = "\x1B[41;1m" ++ s ++ "\x1B[0m"

    pe :: Poss
    pe = listPoss []

    printposs :: Poss -> Poss -> String
    printposs (Poss m1) (Poss m2) =
        let s = flip concatMap [0..8] $ \i -> case (m1 !! i, m2 !! i) of
                    (True, True) -> show (i + 1)
                    (True, False) -> high "."
                    (False, True) -> high (show (i + 1))
                    (False, False) -> " "
        in "(" ++ s ++ ")"


data Action = AResolve Int Value Reason
            | AScratch Int Value Reason
  deriving (Show)

data Reason = RCell | RRow | RCol | RBlock
            | RCellI Int | RRowI Int | RColI Int | RBlockI Int
            | RMultiCellI [Int] [Value]
            | RCombine Reason Reason
  deriving (Show)

type SM = Writer [Action]

reasonReg :: Region -> Reason
reasonReg (Cell i _) = RCellI i
reasonReg (Row r _) = RRowI r
reasonReg (Col c _) = RColI c
reasonReg (Block b _) = RBlockI b

writeAction :: Action -> String
writeAction (AResolve i n reason) =
    "Resolved cell " ++ actionCellString i ++ " to value " ++ show (n + 1) ++
        " (" ++ writeReason reason ++ ")"
writeAction (AScratch i n reason) =
    "Scratched possibility for " ++ show (n + 1) ++ " in cell " ++ actionCellString i ++
        " (" ++ writeReason reason ++ ")"

actionCellString :: Int -> String
actionCellString i = "(" ++ show (i `mod` 9) ++ "," ++ show (i `div` 9) ++ ")"

writeReason :: Reason -> String
writeReason RCell = "cell"
writeReason RRow = "row"
writeReason RCol = "column"
writeReason RBlock = "block"
writeReason (RCellI i) = "situation in cell " ++ show i
writeReason (RRowI i) = "situation in row " ++ show i
writeReason (RColI i) = "situation in column " ++ show i
writeReason (RBlockI i) = "situation in block " ++ show i
writeReason (RMultiCellI idcs vs) =
    "covering of cells " ++ intercalate ", " (map actionCellString idcs) ++
    " with values " ++ intercalate ", " (map (\n -> show (n + 1)) vs)
writeReason (RCombine a b) = writeReason a ++ "; " ++ writeReason b


scratchPossBasicRegion :: Region -> Sudoku
scratchPossBasicRegion reg = maskPoss reg (listPoss $ [0..8] \\ values reg)

scratchAroundCell :: Region -> Sudoku
scratchAroundCell cl@(Cell _ _) =
    let su1 = scratchPossBasicRegion (rowOf cl)
        su2 = scratchPossBasicRegion (colOf (adaptRegion cl su1))
        su3 = scratchPossBasicRegion (blockOf (adaptRegion cl su2))
    in su3

scratchInRegionExcept :: Value -> Region -> Region -> Reason -> SM Sudoku
scratchInRegionExcept val reg except reason =
    foldM (\su i ->
                let cl = cell i (adaptRegion reg su)
                in if cl `cellIsIn` except
                    then return su
                    else if val `possElem` possOf cl
                            then let su' = maskPoss cl (possMakeAllExcept val)
                                 in tell [AScratch (indexOf cl) val reason] >> return su'
                            else return su)
          (sudokuOf reg) (indexRange reg)


tacticPossBasic :: Sudoku -> (Sudoku, [Action])
tacticPossBasic su =
    let su1 = foldl (\su' regc -> scratchPossBasicRegion (regc su'))
                su (concat [[Row i, Col i, Block i] | i <- [0..8]])
    in (su1, [])

scratchPossIndirectReg :: Region -> [Region] -> SM Sudoku
scratchPossIndirectReg mainreg intregs' =
    let intregs = filter (\intreg -> not $ null $ intersectReg mainreg intreg) intregs'
        su = sudokuOf mainreg
        intcells = [intersectReg mainreg intreg | intreg <- intregs]
        alloweds = [foldl possOr emptyPoss (map possOf l) | l <- intcells]
        valOnlyInWhich =  -- [(index, value)] for each value in only one intersection
            concat [maybe [] (\i -> [(i, val)]) (possGetIfInOne alloweds val)
                   | val <- [0..8]]
    in foldM (\su1 (i,v) ->
                scratchInRegionExcept
                    v
                    (adaptRegion (intregs !! i) su1)
                    mainreg
                    (RCombine (reasonReg mainreg) (reasonReg (intregs !! i))))
              su valOnlyInWhich

tacticPossIndirect :: Sudoku -> (Sudoku, [Action])
tacticPossIndirect su =
    let blocks = [Block i su | i <- [0..8]]
        rows = [Row i su | i <- [0..8]]
        cols = [Col i su | i <- [0..8]]
        params = [(Row i, blocks) | i <- [0..8]] ++
                    [(Col i, blocks) | i <- [0..8]] ++
                    [(Block i, rows) | i <- [0..8]] ++
                    [(Block i, cols) | i <- [0..8]]
    in runWriter $
        foldM (\su' (regc, intregs) -> scratchPossIndirectReg (regc su') intregs)
              su params

tacticResolveCells :: Sudoku -> (Sudoku, [Action])
tacticResolveCells su = runWriter $ foldM func su [0..80]
  where
    func :: Sudoku -> Int -> SM Sudoku
    func su' i = case possList (su' `atp` i) of
        [n] -> do
            tell [AResolve i n RCell]
            let su1 = updateValue i n su'
            return $ scratchAroundCell (cell i su1)
        _ -> return su'

tacticResolveRegions :: Sudoku -> (Sudoku, [Action])
tacticResolveRegions su =
    runWriter $ foldM func su [f i | f <- [Row, Col, Block], i <- [0..8]]
  where
    func :: Sudoku -> (Sudoku -> Region) -> SM Sudoku
    func su' regc =
        foldM (\su1 n ->
                    let reg = regc su1
                        ps = map possOf (cells reg)
                    in case possGetIfInOne ps n of
                        Nothing -> return su1
                        Just i -> do
                                    let su_i = indexOf (cell i reg)
                                    tell [AResolve su_i n (reasonReg reg)]
                                    let su2 = updateValue i n reg
                                    return $ scratchAroundCell (cell su_i su2))
              su' [0..8]

tacticPossCover :: Sudoku -> (Sudoku, [Action])
tacticPossCover su =
    runWriter $ foldM func su [f i | f <- [Row, Col, Block], i <- [0..8]]
  where
    func :: Sudoku -> (Sudoku -> Region) -> SM Sudoku
    func su' regc =
        let numsets =
                -- tacticResolveRegions has better Reason for this
                filter (\s -> length s > 1) $
                -- all numbers cover all squares, all right, but not useful
                init $
                subsets ([0..8] \\ values (regc su'))
        in foldM (\su1 ns ->
                    let reg = regc su1
                        ps = map possOf (cells reg)
                        idcslist =
                            [map fst $
                                filter (\(_, p) -> n `possElem` p) $ zip [0..8] ps
                            | n <- ns]
                        idcs = sort $ foldl union [] idcslist
                    in if length idcs == length ns &&
                                -- tacticResolveRegions has better Reason for this
                                all (\l -> length l > 1) idcslist
                        then applyCover reg ns idcs
                        else return su1)
                 su' numsets

    applyCover :: Region -> [Value] -> [Int] -> SM Sudoku
    applyCover reg ns idcs =
        let su' = sudokuOf reg
            mask = listPoss ns
            reasonIdcs = [indexOf (cell i reg) | i <- idcs]
            reason = RCombine (reasonReg reg) (RMultiCellI reasonIdcs ns)
        in foldM (\su1 i ->
                    let reg1 = adaptRegion reg su1
                    in do
                        forM_ (possList (possOf (cell i reg1)) \\ ns) $
                              \n -> tell [AScratch (indexOf (cell i reg1)) n reason]
                        return $ maskPoss (cell i reg1) mask)
                 su' idcs


operations :: [Sudoku -> (Sudoku, [Action])]
operations = [
    tacticPossBasic, tacticPossIndirect, tacticPossCover,
    tacticResolveCells, tacticResolveRegions]

solve :: Sudoku -> [(Sudoku, [Action])]
solve su =
    let l = tail $ go su operations
        (su', _) = last l
    in if su' == su then [] else l ++ solve su'
  where
    go :: Sudoku -> [Sudoku -> (Sudoku, [Action])] -> [(Sudoku, [Action])]
    go su1 ops = scanl (\(su2, _) op -> op su2) (su1, []) ops

main :: IO ()
main = do
    su <- liftM readSudoku getContents

    let useBT = False
        btSols = backtrackSolve su

    when useBT $
        case btSols of
            [] -> do
                die "Cannot solve sudoku using backtracking"
            list -> do
                putStrLn $ "Backtracking gave " ++ show (length list) ++ " solutions:"
                forM_ list $ \res -> do
                    putStrLn $ writeSudoku res
                    putStr "\n"

    let reslist = (su, []) : solve su
        suRes = fst (last reslist)

    putStrLn (writeSudokuFull su)

    mapM_ (\((su', acts), (prevsu, _)) -> do
                putStr "\n\n"
                mapM_ (putStrLn . writeAction) acts
                -- print acts
                putStrLn (writeSudokuDiff prevsu su'))
           (zip (tail reslist) reslist)

    putStrLn (writeSudoku suRes)
    putStr "\n"

    when useBT $
        when (length btSols /= 1 || suRes /= btSols !! 0) $
            die "Incorrectly solved!"