module Control.Monad
    ( module Control.Applicative
    , Monad (..)
    , (=<<)
    , (<=<)
    , bind
    , (>>=)
    , (>=>)
    , extend2
    , extend3
    , extend4
    , extend5
    , Kleisli (..)
    , Comonad (..)
    , (<<=)
    , (=<=)
    , cobind
    , (=>>)
    , (=>=)
    , CoKleisli (..)
    ) where

import GHC.IOBase

import Control.Applicative
import Data.Function
import Data.Wrapper

class Applicative (~>) m => Monad (~>) m where

    extend :: (a ~> m b) -> (m a ~> m b)
    extend k = join . map k

    join :: m (m a) ~> m a
    join = extend id

instance Monad Hask (Hask a) where
    join f x = f x x

instance Monad Hask IO where
    extend = flip bindIO

infixr 1 =<<
(=<<) :: Monad (~>) m => (a ~> m b) -> (m a ~> m b)
(=<<) = extend

infixr 1 <=<
(<=<) :: Monad (~>) m => (b ~> m c) -> (a ~> m b) -> (a ~> m c)
(<=<) = result . extend

bind :: Monad Hask m => m a -> (a -> m b) -> m b
bind = flip extend

infixl 1 >>=
(>>=) :: Monad Hask m => m a -> (a -> m b) -> m b
(>>=) = bind

infixr 1 >=>
(>=>) :: Monad (~>) m => (a ~> m b) -> (b ~> m c) -> (a ~> m c)
(>=>) = flip (<=<)

extend2 :: Monad Hask m => (a -> b -> m c) -> (m a -> m b -> m c)
extend2 = result extend . ignoreArg2 extend

extend3 :: Monad Hask m => (a -> b -> c -> m d) -> (m a -> m b -> m c -> m d)
extend3 = result2 extend . ignoreArg3 extend2

extend4 :: Monad Hask m  =>
           (a -> b -> c -> d -> m e) -> (m a -> m b -> m c -> m d -> m e)
extend4 = result3 extend . ignoreArg4 extend3

extend5 :: Monad Hask m =>
           (a -> b -> c -> d -> e -> m f) -> (m a -> m b -> m c -> m d -> m e -> m f)
extend5 = result4 extend . ignoreArg5 extend4

newtype Kleisli (~>) m a b = Kleisli { runKleisli :: a ~> m b }

instance Wrapper (Kleisli (~>) m a b) where
    type Inner (Kleisli (~>) m a b) = a ~> m b
    wrap = Kleisli
    unwrap = runKleisli

instance Monad (~>) m => Category (Kleisli (~>) m) where
    id = Kleisli point
    (.) = inWrapper2 (<=<)

class Copointed (~>) w => Comonad (~>) w where
    coextend :: (w a ~> b) -> (w a ~> w b)
    coextend f = map f . cojoin

    cojoin :: w a ~> w (w a)
    cojoin = coextend id

(<<=) :: Comonad (~>) w => (w a ~> b) -> (w a ~> w b)
(<<=) = coextend

(=<=) :: Comonad (~>) w => (w b ~> c) -> (w a ~> b) -> (w a ~> c)
(=<=) f = result f . coextend

cobind :: Comonad Hask w => w a -> (w a -> b) -> w b
cobind = flip coextend

(=>>) :: Comonad Hask w => w a -> (w a -> b) -> w b
(=>>) = cobind

(=>=) :: Comonad (~>) w => (w a ~> b) -> (w b ~> c) -> (w a ~> c)
(=>=) = flip (=<=)

newtype CoKleisli (~>) w a b = CoKleisli { runCoKleisli :: w a ~> b }

instance Wrapper (CoKleisli (~>) w a b) where
    type Inner (CoKleisli (~>) w a b) = w a ~> b
    wrap = CoKleisli
    unwrap = runCoKleisli

instance Comonad (~>) w => Category (CoKleisli (~>) w) where
    id = CoKleisli extract
    (.) = inWrapper2 (=<=)