module Problem164 where

import Data.Array
import Test.QuickCheck

cnts :: Array (Int, Int, Int, Int) Int
-- (n,b,c) -> count of n-digit numbers with starting b-digit sum == c and 
-- any 3 consecutive digit sum <= 9
cnts = array
    ((1, 1, 0, 0), (20, 2, 9, 9))
    [ ((l, c, s, k), f l c s k) | l <- [1 .. 20], c <- [1, 2], s <- [0 .. 9], k<-[0..9] ]
  where
    f l c s  k
        | l <= 2 = brute (l, c, s, k)
        | c == 1 && s /= k  = 0
        | s < k = 0
        | otherwise = f' l c s k
    -- f 4 1 1 = sum [cnts ! (3, 2, s) | s<- [0..8]]
    -- A|BC...
    f' l 1 s _ = sum [cnts ! (l-1, 2, s', k') | k'<- [0..9], s'<-[0..9-s]]
    -- AB|CD...
    -- s=k+b, s+c <=9 && b+c+d <=9 => c <= 9-s && s' = c+d <= 9-b = 9-(s-k)
    f' l 2 s k = sum [ cnts ! (l-2, 2,s', k') | k' <-[0..9-s], s'<-[0..9-s+k]]

newtype Key = Key (Int, Int, Int, Int) deriving (Show)

instance Arbitrary Key where
    arbitrary = Key <$> ((,,,) <$> choose (1,5) <*> choose (1, 2) <*> choose (0, 9) <*> choose (0,9))

prop_cnt :: Key -> Bool
prop_cnt (Key k) = cnts ! k == brute k

brute :: (Int, Int, Int, Int) -> Int
brute = length . brute'

brute' :: (Int, Int, Int, Int) -> [[Int]]
brute' (n, b, c, k) =
    filter (\xs -> valid xs && sum (take b xs) == c && head xs == k) $ mapM (const [0 .. 9]) [1 .. n]

valid :: (Ord a, Num a) => [a] -> Bool
valid xs = all (<= 9) $ zipWith3 (\a b c -> a + b + c) xs (tail xs) (drop 2 xs)

main :: IO ()
main = print $ sum [cnts ! (20, 1, k, k) | k <-[1..9]]
