import Data.List (group)
import Control.Monad (forM_)

replicateAtLeastOne :: Int -> a -> [a]
replicateAtLeastOne n x = x : replicate (n-1) x

zipWithLazy :: (a -> b -> c) -> [a] -> [b] -> [c]
zipWithLazy f ~(x:xs) ~(y:ys) = f x y : zipWithLazy f xs ys

kolakoski :: [Int] -> [Int]
kolakoski items = s
  where s = concat $ zipWithLazy replicateAtLeastOne s $ cycle items

rle :: Eq a => [a] -> [Int]
rle = map length . group

sameAsRleUpTo :: Int -> [Int] -> Bool
sameAsRleUpTo n s = r == take (length r) prefix
  where prefix = take n s
        r = init $ rle prefix

main :: IO ()
main = forM_ [([1, 2], 20),
              ([2, 1], 20),
              ([1, 3, 1, 2], 30),
              ([1, 3, 2, 1], 30)]
        $ \(items, n) -> do
          putStrLn $ "First " ++ show n ++ " members of the sequence generated by " ++ show items ++ ":"
          let s = kolakoski items
          print $ take n s
          putStrLn $ "Possible Kolakoski sequence? " ++ show (sameAsRleUpTo n s)
          putStrLn ""