import Control.Applicative (liftA2, ZipList(..))
import Control.Applicative.Backwards
import Data.Foldable (toList, traverse_)
import Data.Maybe (fromMaybe)


class Applicative f => Alternative f where
  empty :: f a
  (<|>) :: f a -> f a -> f a

  some :: f a -> f [a]
  some v = some_v
    where
      many_v = fromMaybe [] <$> optional some_v
      some_v = liftA2 (:) v many_v

  many :: f a -> f [a]
  many v = many_v
    where
      many_v = fromMaybe [] <$> optional some_v
      some_v = liftA2 (:) v many_v

  -- new!
  optional :: f a -> f (Maybe a)
  optional v = Just <$> v <|> pure Nothing

infixl 3 <|>


instance Alternative Maybe where
  empty = Nothing
  Nothing <|> r = r
  l       <|> _ = l

  -- new!
  optional = Just

instance Alternative [] where
  empty = []
  (<|>) = (++)

  -- new!
  optional xs = let
    (h, t) = case xs of
      x : xs' -> (Just x, optional xs')
      []      -> (Nothing, [])
    in h : t

instance Alternative ZipList where
  empty = ZipList []
  ZipList xs0 <|> ZipList ys0 = ZipList $ go xs0 ys0
    where
      go (x:xs) (_:ys) = x : go xs ys
      go    []     ys  = ys

  -- new!
  optional (ZipList xs) = let
    (h, ZipList t) = case xs of
      x : xs' -> (Just x, optional $ ZipList xs')
      []      -> (Nothing, pure Nothing)
    in ZipList $ h : t


instance Alternative f => Alternative (Backwards f) where
  empty = Backwards empty
  Backwards x <|> Backwards y = Backwards (x <|> y)
  optional (Backwards x) = Backwards (optional x)


-- | Create an infinite list of infinite lists with elements
-- drawn from the input list. I am bad at naming functions.
--
-- >>> traverse_ print $ take 10 $ map (take 4) $ wezzle [0..2]
-- [0,0,0,0]
-- [1,0,0,0]
-- [2,0,0,0]
-- [0,1,0,0]
-- [1,1,0,0]
-- [2,1,0,0]
-- [0,2,0,0]
-- [1,2,0,0]
-- [2,2,0,0]
-- [0,0,1,0]
wezzle :: [a] -> [[a]]
wezzle = forwards . some . Backwards


newtype Lazy f a = Lazy { greedy :: f a }
  deriving (Functor, Applicative)

instance Alternative f => Alternative (Lazy f) where
  empty = Lazy empty
  Lazy l <|> Lazy r = Lazy $ r <|> l


-- | Create an infinite list of finite lists of every length,
-- with elements drawn from the input list. I am still bad at
-- naming functions.
--
-- >>> traverse_ print $ take 10 $ staggle [0,1]
-- []
-- [0]
-- [1]
-- [0,0]
-- [0,1]
-- [1,0]
-- [1,1]
-- [0,0,0]
-- [0,0,1]
-- [0,1,0]
staggle :: [a] -> [[a]]
staggle = greedy . forwards . many . Backwards . Lazy


main :: IO ()
main = do
  print $ fmap (take 4) $ some $ Just 0
  print $ fmap (take 4) $ many $ Just 0
  print $ take 5 $ map (take 4) $ some [0..2]
  print $ take 5 $ map (take 4) $ many [0..2]
  print $ take 5 $ map (take 4) $ getZipList $ some $ ZipList [0..2]
  print $ take 5 $ map (take 4) $ getZipList $ many $ ZipList [0..2]
  traverse_ print $ take 10 $ map (take 4) $ wezzle [0..2]
  traverse_ print $ take 10 $ staggle [0,1]