-- A simple 'level' monad type

module MonadDiags where

import Control.Applicative
import Control.Monad

apD :: [[a -> b]] -> [[a]] -> [[b]]
apD [] _  = []
apD _  [] = []
apD (xx:xs) ys = unionD (map apXX ys) ([] : apD xs ys)
   where apXX yy = xx <*> yy


unionD :: [[a]] -> [[a]] -> [[a]]
unionD [] ys = ys
unionD xs [] = xs
unionD (x:xs) (y:ys) = (x++y) : unionD xs ys


-- Now to wrap things up in an applicative functor

newtype Diags a = Diags { runDiags :: [[a]] }

instance Functor Diags where
   fmap f = Diags . map (map f) . runDiags

instance Applicative Diags where
   pure a = Diags [[a]]
   a <*> b = Diags (runDiags a `apD` runDiags b)

instance Alternative Diags where
   empty = Diags [[]]
   a <|> b = Diags (runDiags a `unionD` runDiags b)

each :: [a] -> Diags a
each = Diags . map return


-- And if we want a monad, we should also have a join function

joinD :: [[Diags a]] -> [[a]]
joinD [] = []
joinD (xx:xs) = unionsD (map runDiags xx) `unionD` ([] : joinD xs)

unionsD :: [[[a]]] -> [[a]]
unionsD = foldr unionD []

instance Monad Diags where
    return = pure
    a >>= b = Diags (joinD (runDiags $ fmap b a))
    fail _ = empty

instance MonadPlus Diags where
    mzero = empty
    mplus = (<|>)