{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TupleSections         #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE TypeFamilies          #-}
module Data.Profunctor.Optic.Operator (
    type (+)
  , (&)
  , rgt
  , rgt'
  , lft
  , lft'
  , swap
  , eswap
  , fork
  , join
  , eval
  , apply
  , branch
  , branch'
  , assocl
  , assocr
  , assocl'
  , assocr'
  , eassocl
  , eassocr
  , parr
  , unarr
  , peval
  , constl
  , constr
  , shiftl
  , shiftr
  , coercel
  , coercer
  , coercel'
  , coercer'
  , strong
  , costrong
  , choice
  , cochoice
  , pull
  , repn
  , corepn
  , star
  , costar
  , unstar
  , uncostar
  , sieve'
  , tabulate'
  , cosieve'
  , cotabulate'
  , pushr
  , pushl
  , liftR2
  , pdivide
  , pappend
  , (<<*>>)
  , (****)
  , (&&&&)
) where

import Data.Function
import Data.Profunctor.Closed
import Data.Profunctor.Optic.Types
import Data.Profunctor.Optic.Import

branch :: (a -> Bool) -> b -> c -> a -> b + c
branch f y z x = if f x then Right z else Left y
{-# INLINE branch #-}

branch' :: (a -> Bool) -> a -> a + a
branch' f x = branch f x x x
{-# INLINE branch' #-}

assocl :: (a , (b , c)) -> ((a , b) , c)
assocl (a, (b, c)) = ((a, b), c)
{-# INLINE assocl #-}

assocr :: ((a , b) , c) -> (a , (b , c))
assocr ((a, b), c) = (a, (b, c))
{-# INLINE assocr #-}

assocl' :: (a , b + c) -> (a , b) + c
assocl' = eswap . traverse eswap
{-# INLINE assocl' #-}

assocr' :: (a + b , c) -> a + (b , c)
assocr' (f, b) = fmap (,b) f
{-# INLINE assocr' #-}

eassocl :: a + (b + c) -> (a + b) + c
eassocl (Left a)          = Left (Left a)
eassocl (Right (Left b))  = Left (Right b)
eassocl (Right (Right c)) = Right c
{-# INLINE eassocl #-}

eassocr :: (a + b) + c -> a + (b + c)
eassocr (Left (Left a))  = Left a
eassocr (Left (Right b)) = Right (Left b)
eassocr (Right c)        = Right (Right c)
{-# INLINE eassocr #-}

{-
eassocr' :: (a -> b) + c -> a -> b + c
eassocr' abc a = either (\ab -> Left $ ab a) Right abc
{-# INLINE eassocr' #-}

forget1 :: ((c, a) -> (c, b)) -> a -> b
forget1 f a = b where (c, b) = f (c, a)
{-# INLINE forget1 #-}

forget2 :: ((a, c) -> (b, c)) -> a -> b
forget2 f a = b where (b, c) = f (a, c)
{-# INLINE forget2 #-}

forgetl :: (c + a -> c + b) -> a -> b
forgetl f = go . Right where go = either (go . Left) id . f
{-# INLINE forgetl #-}

forgetr :: (a + c -> b + c) -> a -> b
forgetr f = go . Left where go = either id (go . Right) . f
{-# INLINE forgetr #-}
-}

parr :: Representable p => Applicative (Rep p) => (a -> b) -> p a b
parr = tabulate . (pure .)
{-# INLINE parr #-}

unarr :: Coapplicative w => Sieve p w => p a b -> a -> b
unarr = (copure .) . sieve
{-# INLINE unarr #-}

peval :: Strong p => p a (a -> b) -> p a b
peval = rmap eval . pull
{-# INLINE peval #-}

constl :: Profunctor p => b -> p b c -> p a c
constl = lmap . const
{-# INLINE constl #-}

constr :: Profunctor p => c -> p a b -> p a c
constr = rmap . const
{-# INLINE constr #-}

shiftl :: Profunctor p => p (a + b) c -> p b (c + d)
shiftl = dimap Right Left
{-# INLINE shiftl #-}

shiftr :: Profunctor p => p b (c , d) -> p (a , b) c
shiftr = dimap snd fst
{-# INLINE shiftr #-}

coercel :: Profunctor p => Bifunctor p => p a b -> p c b
coercel = first absurd . lmap absurd
{-# INLINE coercel #-}

coercer :: Profunctor p => Contravariant (p a) => p a b -> p a c
coercer = rmap absurd . contramap absurd
{-# INLINE coercer #-}

coercel' :: Corepresentable p => Contravariant (Corep p) => p a b -> p c b
coercel' = corepn (. phantom)
{-# INLINE coercel' #-}

coercer' :: Representable p => Contravariant (Rep p) => p a b -> p a c
coercer' = repn (phantom .)
{-# INLINE coercer' #-}

strong :: Strong p => ((a , b) -> c) -> p a b -> p a c
strong f = dimap fork f . second'
{-# INLINE strong #-}

costrong :: Costrong p => ((a , b) -> c) -> p c a -> p b a
costrong f = unsecond . dimap f fork
{-# INLINE costrong #-}

choice :: Choice p => (c -> (a + b)) -> p b a -> p c a
choice f = dimap f join . right'
{-# INLINE choice #-}

cochoice :: Cochoice p => (c -> (a + b)) -> p a c -> p a b
cochoice f = unright . dimap join f
{-# INLINE cochoice #-}

pull :: Strong p => p a b -> p a (a , b)
pull = lmap fork . second'
{-# INLINE pull #-}

repn :: Representable p => ((a -> Rep p b) -> s -> Rep p t) -> p a b -> p s t
repn f = tabulate . f . sieve
{-# INLINE repn #-}

corepn :: Corepresentable p => ((Corep p a -> b) -> Corep p s -> t) -> p a b -> p s t
corepn f = cotabulate . f . cosieve
{-# INLINE corepn #-}

star :: Applicative f => Star f a a
star = Star pure
{-# INLINE star #-}

costar :: Coapplicative f => Costar f a a
costar = Costar copure
{-# INLINE costar #-}

unstar :: Coapplicative f => Star f a b -> a -> b
unstar f = copure . runStar f
{-# INLINE unstar #-}

uncostar :: Applicative f => Costar f a b -> a -> b
uncostar f = runCostar f . pure
{-# INLINE uncostar #-}

sieve' :: Sieve p f => p d c -> Star f d c
sieve' = Star . sieve
{-# INLINE sieve' #-}

tabulate' :: Representable p => Star (Rep p) a b -> p a b
tabulate' = tabulate . runStar
{-# INLINE tabulate' #-}

cosieve' :: Cosieve p f => p a b -> Costar f a b
cosieve' = Costar . cosieve
{-# INLINE cosieve' #-}

cotabulate' :: Corepresentable p => Costar (Corep p) a b -> p a b
cotabulate' = cotabulate . runCostar
{-# INLINE cotabulate' #-}

pushr :: Closed p => Representable p => Apply (Rep p) => p (a , b) c -> p a b -> p a c
pushr = (<<*>>) . curry'
{-# INLINE pushr #-}

pushl :: Closed p => Representable p => Apply (Rep p) => p a c -> p b c -> p a (b -> c)
pushl p q = curry' $ pdivide id p q
{-# INLINE pushl #-}

{-

λ> :t foo
foo
  :: (Corepresentable p1, Coapplicative (Corep p1),
      Cosieve p2 (Corep p1)) =>
     p2 a c -> p1 (c + a) c

pliftC
  :: (Corepresentable p1, Coapplicative (Corep p1),
      Cosieve p2 (Either b)) =>
     p2 a c -> p1 (b + a) c
pliftC x = cotabulate $ cosieve x . bimap copure copure . coapply


pliftC
  :: (Corepresentable p1, Coapply (Corep p1),
      Cosieve p2 (Either (Corep p1 a))) =>
     p2 (Corep p1 b) c -> p1 (a + b) c
--pliftC :: p (a + b) c -> (p a c + p b c)
pliftC x = cotabulate $ cosieve x . coapply
pliftC
  :: (Corepresentable p, Coapply (Corep p), Cosieve p f) =>
     (Corep p a -> f a2)
     -> (Corep p b -> f a2) -> p a2 c -> p (a + b) c
pliftC f g x = cotabulate $ cosieve x . either f g . coapply

--TODO: 
-- add Data.Profunctor.Fold w/ Mealy/Moore types
-- add coapplicative profunctor operators

liftR2 :: Representable p => Applicative (Rep p) => (b -> c -> d) -> p a b -> p a c -> p a d
liftR2 f x y = tabulate $ \s -> liftA2 f (sieve x s) (sieve y s)
{-# INLINE liftR2 #-}
-}


liftR2 :: Representable p => Apply (Rep p) => (b -> c -> d) -> p a b -> p a c -> p a d
liftR2 f x y = tabulate $ \s -> liftF2 f (sieve x s) (sieve y s)
{-# INLINE liftR2 #-}

infixr 2 ||||

(||||) :: Corepresentable p => Coapply (Corep p) => p a1 b -> p a2 b -> p (a1 + a2) b
p |||| q = cotabulate $ either (cosieve p) (cosieve q) . coapply
{-# INLINE (||||) #-}

infixr 3 &&&&

(&&&&) ::  Representable p => Apply (Rep p) => p a b1 -> p a b2 -> p a (b1 , b2)
p &&&& q = liftR2 (,) p q
{-# INLINE (&&&&) #-}

infixl 4 <<*>>

(<<*>>) :: Representable p => Apply (Rep p) => p a (b -> c) -> p a b -> p a c
(<<*>>) = liftR2 ($)
{-# INLINE (<<*>>) #-}

infixr 3 ****

(****) :: Representable p => Apply (Rep p) => p a1 b1 -> p a2 b2 -> p (a1 , a2) (b1 , b2)
p **** q = dimap fst (,) p <<*>> lmap snd q
{-# INLINE (****) #-}

pdivide :: Representable p => Apply (Rep p) => (a -> (a1 , a2)) -> p a1 b -> p a2 b -> p a b
pdivide f p q = dimap f fst $ dimap fst (,) p <<*>> lmap snd q
{-# INLINE pdivide #-}

pappend :: Representable p => Apply (Rep p) => p a b -> p a b -> p a b
pappend = pdivide fork
{-# INLINE pappend #-}