{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE LambdaCase #-}

module Data.Op2 where

import Control.Applicative (Applicative(pure, (<*>), liftA2), Alternative(..))
import Control.Category ( Category(..) )
import Control.Lens
    ( view,
      iso,
      prism',
      (#),
      review,
      _Wrapped,
      Identity(..),
      Iso,
      Lens',
      Prism',
      Rewrapped,
      Wrapped(..) )
import Control.Monad ( Monad(return, (>>=)) )
import Control.Monad.Trans( MonadTrans(..) )
import Control.Monad.IO.Class ( MonadIO(..) )
import Data.Bool ( Bool(..) )
import Data.Either ( Either(Left, Right), either )
import Data.Eq ( Eq((==)) )
import Data.Functor( Functor(..), (<$>) )
import Data.Functor.Apply( Apply((<.>)) )
import Data.Functor.Alt( Alt((<!>)) )
import Data.Functor.Bind( Bind((>>-)) )
import Data.Functor.Contravariant(Contravariant(contramap), Equivalence(Equivalence), Comparison(Comparison))
import Data.Functor.Contravariant.Divisible(Divisible(divide, conquer))
import Data.Foldable ( Foldable )
import qualified Data.List as List(sortBy)
import Data.List.NonEmpty(NonEmpty, groupBy)
import qualified Data.List.NonEmpty as NonEmpty(sortBy)
import Data.Maybe ( Maybe(Nothing, Just) )
import Data.Monoid(Monoid(mempty), All, Any)
import Data.Ord ( Ord(compare), Ordering(..) )
import Data.Profunctor( Profunctor(dimap), Choice(..) )
import Data.Semigroup ( Semigroup((<>)) )

newtype OpT2 b f a =
  OpT2 (a -> a -> f b)

type Op2 b a =
  OpT2 b Identity a

op2 ::
  Iso
    (Op2 b a)
    (Op2 b a)
    (a -> a -> b)
    (a -> a -> b)
op2 =
  iso
    (\o a1 a2 -> runIdentity (view _Wrapped o a1 a2))
    (\f -> OpT2 (\a1 a2 -> Identity (f a1 a2)))

instance OpT2 b' f' a' ~ t =>
  Rewrapped (OpT2 b f a) t

instance Wrapped (OpT2 b f a) where
  type Unwrapped (OpT2 b f a) = a -> a -> f b
  _Wrapped' = iso (\(OpT2 x) -> x) OpT2

instance (Applicative f, Semigroup b) => Semigroup (OpT2 b f a) where
  OpT2 f <> OpT2 g =
    OpT2 (\a1 a2 -> liftA2 (<>) (f a1 a2) (g a1 a2))

instance (Applicative f, Monoid b) => Monoid (OpT2 b f a) where
  mempty =
    OpT2 (\_ _ -> pure mempty)

instance Contravariant (OpT2 b f) where
  contramap f (OpT2 g) =
    OpT2 (\b1 b2 -> g (f b1) (f b2))

instance (Applicative f, Monoid b) => Divisible (OpT2 b f) where
  divide f (OpT2 g) (OpT2 h) =
    OpT2 (\a1 a2 ->
          let (b1, c1) = f a1
              (b2, c2) = f a2
          in  liftA2 (<>) (g b1 b2) (h c1 c2))
  conquer =
    mempty

newtype MonadOpT2 a f b =
  MonadOpT2 (OpT2 b f a)

monadOp2 ::
  Iso
    (OpT2 b f a)
    (OpT2 b' f' a')
    (MonadOpT2 a f b)
    (MonadOpT2 a' f' b')
monadOp2 =
  iso
    MonadOpT2
    (\(MonadOpT2 x) -> x)

argument1 ::
  Applicative f =>
  MonadOpT2 a f a
argument1 =
  MonadOpT2 (OpT2 (\a1 _ -> pure a1))

argument2 ::
  Applicative f =>
  MonadOpT2 a f a
argument2 =
  MonadOpT2 (OpT2 (\_ a2 -> pure a2))

instance Functor f => Functor (MonadOpT2 a f) where
  fmap f (MonadOpT2 (OpT2 g)) =
    MonadOpT2 (OpT2 (\a1 a2 -> fmap f (g a1 a2)))

instance Apply f => Apply (MonadOpT2 a f) where
  MonadOpT2 (OpT2 f) <.> MonadOpT2 (OpT2 g) =
    MonadOpT2 (OpT2 (\a1 a2 -> f a1 a2 <.> g a1 a2))

instance Applicative f => Applicative (MonadOpT2 a f) where
  pure a =
    MonadOpT2 (OpT2 (\_ _ -> pure a))
  MonadOpT2 (OpT2 f) <*> MonadOpT2 (OpT2 g) =
    MonadOpT2 (OpT2 (\a1 a2 -> f a1 a2 <*> g a1 a2))

instance Bind f => Bind (MonadOpT2 a f) where
  MonadOpT2 (OpT2 f) >>- g =
    MonadOpT2 (OpT2 (\a1 a2 -> f a1 a2 >>- \a -> let MonadOpT2 (OpT2 k) = g a in k a1 a2))

instance Monad f => Monad (MonadOpT2 a f) where
  return a =
    MonadOpT2 (OpT2 (\_ _ -> return a))
  MonadOpT2 (OpT2 f) >>= g =
    MonadOpT2 (OpT2 (\a1 a2 -> f a1 a2 >>= \a -> let MonadOpT2 (OpT2 k) = g a in k a1 a2))

instance Alt f => Alt (MonadOpT2 a f) where
  MonadOpT2 (OpT2 f) <!> MonadOpT2 (OpT2 g) =
    MonadOpT2 (OpT2 (\a1 a2 -> f a1 a2 <!> g a1 a2))

instance Alternative f => Alternative (MonadOpT2 a f) where
  empty =
    MonadOpT2 (OpT2 (\_ _ -> empty))
  MonadOpT2 (OpT2 f) <|> MonadOpT2 (OpT2 g) =
    MonadOpT2 (OpT2 (\a1 a2 -> f a1 a2 <|> g a1 a2))

instance MonadIO f => MonadIO (MonadOpT2 a f) where
  liftIO a =
    MonadOpT2 (OpT2 (\ _ _ -> liftIO a))

instance MonadTrans (MonadOpT2 a) where
  lift a =
    MonadOpT2 (OpT2 (\_ _ -> a))

newtype ProfunctorOpT2 f a b =
  ProfunctorOpT2 (OpT2 b f a)

profunctorOp2 ::
  Iso
    (OpT2 b f a)
    (OpT2 b' f' a')
    (ProfunctorOpT2 f a b)
    (ProfunctorOpT2 f' a' b')
profunctorOp2 =
  iso
    ProfunctorOpT2
    (\(ProfunctorOpT2 x) -> x)

instance Functor f => Profunctor (ProfunctorOpT2 f) where
  dimap f g (ProfunctorOpT2 (OpT2 x)) =
    ProfunctorOpT2 (OpT2 (\a1 a2 -> fmap g (x (f a1) (f a2))))

instance Applicative f => Choice (ProfunctorOpT2 f) where
  left' (ProfunctorOpT2 (OpT2 x))  =
    let lft l = either l (pure . Right)
    in  ProfunctorOpT2 (OpT2 (\a1 a2 -> lft (\a1' -> lft (fmap Left . x a1') a2) a1))
  right' (ProfunctorOpT2 (OpT2 x))  =
    let rgt r = either (pure . Left) r
    in  ProfunctorOpT2 (OpT2 (\a1 a2 -> rgt (\a1' -> rgt (fmap Right . x a1') a2) a1))

instance Functor f => Functor (ProfunctorOpT2 f a) where
  fmap f (ProfunctorOpT2 (OpT2 g)) =
    ProfunctorOpT2 (OpT2 (\a1 a2 -> fmap f (g a1 a2)))

instance Apply f => Apply (ProfunctorOpT2 f a) where
  ProfunctorOpT2 (OpT2 f) <.> ProfunctorOpT2 (OpT2 g) =
    ProfunctorOpT2 (OpT2 (\a1 a2 -> f a1 a2 <.> g a1 a2))

instance Applicative f => Applicative (ProfunctorOpT2 f a) where
  pure a =
    ProfunctorOpT2 (OpT2 (\_ _ -> pure a))
  ProfunctorOpT2 (OpT2 f) <*> ProfunctorOpT2 (OpT2 g) =
    ProfunctorOpT2 (OpT2 (\a1 a2 -> f a1 a2 <*> g a1 a2))

instance Bind f => Bind (ProfunctorOpT2 f a) where
  ProfunctorOpT2 (OpT2 f) >>- g =
    ProfunctorOpT2 (OpT2 (\a1 a2 -> f a1 a2 >>- \a -> let ProfunctorOpT2 (OpT2 k) = g a in k a1 a2))

instance Monad f => Monad (ProfunctorOpT2 f a) where
  return a =
    ProfunctorOpT2 (OpT2 (\_ _ -> return a))
  ProfunctorOpT2 (OpT2 f) >>= g =
    ProfunctorOpT2 (OpT2 (\a1 a2 -> f a1 a2 >>= \a -> let ProfunctorOpT2 (OpT2 k) = g a in k a1 a2))

instance Alt f => Alt (ProfunctorOpT2 f a) where
  ProfunctorOpT2 (OpT2 f) <!> ProfunctorOpT2 (OpT2 g) =
    ProfunctorOpT2 (OpT2 (\a1 a2 -> f a1 a2 <!> g a1 a2))

instance Alternative f => Alternative (ProfunctorOpT2 f a) where
  empty =
    ProfunctorOpT2 (OpT2 (\_ _ -> empty))
  ProfunctorOpT2 (OpT2 f) <|> ProfunctorOpT2 (OpT2 g) =
    ProfunctorOpT2 (OpT2 (\a1 a2 -> f a1 a2 <|> g a1 a2))

instance MonadIO f => MonadIO (ProfunctorOpT2 f a) where
  liftIO a =
    ProfunctorOpT2 (OpT2 (\ _ _ -> liftIO a))

bothOp2T ::
  (Applicative f, Semigroup b) =>
  (a -> f b)
  -> OpT2 b f a
bothOp2T f =
  OpT2 (\a1 a2 -> (<>) <$> f a1 <*> f a2)

bothOp2 ::
  Semigroup b =>
  (a -> b)
  -> Op2 b a
bothOp2 f =
  bothOp2T (Identity . f)

appendOp2 ::
  (Applicative f, Semigroup x) =>
  OpT2 x f x
appendOp2 =
  bothOp2T pure

class AsBool a where
  _Bool ::
    Prism' a Bool
  _True ::
    Prism' a ()
  _True =
    _Bool .
    prism'
      (\() -> True)
      (\case  True ->
                Just ()
              _ ->
                Nothing)
  _False ::
    Prism' a ()
  _False =
    _Bool .
    prism'
      (\() -> False)
      (\case  False ->
                Just ()
              _ ->
                Nothing)

instance AsBool Bool where
  _Bool =
    id

instance AsBool All where
  _Bool =
    _Wrapped

instance AsBool Any where
  _Bool =
    _Wrapped

class HasBool a where
  bool' ::
    Lens' a Bool

instance HasBool Bool where
  bool' =
    id

instance HasBool All where
  bool' =
    _Wrapped

instance HasBool Any where
  bool' =
    _Wrapped

true ::
  AsBool a =>
  a
true =
  _True # ()

false ::
  AsBool a =>
  a
false =
  _False # ()

class AsOrdering a where
  _Ordering ::
    Prism' a Ordering
  _LT ::
    Prism' a ()
  _LT =
    _Ordering .
    prism'
      (\() -> LT)
      (\case  LT ->
                Just ()
              _ ->
                Nothing
      )
  _EQ ::
    Prism' a ()
  _EQ =
    _Ordering .
    prism'
      (\() -> EQ)
      (\case  EQ ->
                Just ()
              _ ->
                Nothing
      )
  _GT ::
    Prism' a ()
  _GT =
    _Ordering .
    prism'
      (\() -> GT)
      (\case  GT ->
                Just ()
              _ ->
                Nothing
      )

instance AsOrdering Ordering where
  _Ordering =
    id

class HasOrdering a where
  ordering ::
    Lens' a Ordering

instance HasOrdering Ordering where
  ordering =
    id

lt ::
  AsOrdering a =>
  a
lt =
  _LT # ()

eq ::
  AsOrdering a =>
  a
eq =
  _EQ # ()

gt ::
  AsOrdering a =>
  a
gt =
  _GT # ()

equivalence ::
  Iso
    (Equivalence a)
    (Equivalence a')
    (Op2 Bool a)
    (Op2 Bool a')
equivalence =
  iso
    (\(Equivalence x) -> review op2 x)
    (Equivalence . view op2)

comparison ::
  Iso
    (Comparison a)
    (Comparison a')
    (Op2 Ordering a)
    (Op2 Ordering a')
comparison =
  iso
    (\(Comparison x) -> review op2 x)
    (Comparison . view op2)

eqOp2 ::
  (Eq a, AsBool b) =>
  Op2 b a
eqOp2 =
  review op2 (\a1 a2 -> review _Bool (a1 == a2))

ordOp2 ::
  (Ord a, AsOrdering b) =>
  Op2 b a
ordOp2 =
  review op2 (\a1 a2 -> review _Ordering (a1 `compare` a2))

groupBy' ::
  (HasBool b, Foldable f) =>
  Op2 b a
  -> f a
  -> [NonEmpty a]
groupBy' f =
  groupBy (\a1 a2 -> view bool' (view op2 f a1 a2))

sortBy ::
  HasOrdering b =>
  Op2 b a
  -> [a]
  -> [a]
sortBy f =
  List.sortBy (\a1 a2 -> view ordering (view op2 f a1 a2))

sortBy1 ::
  HasOrdering b =>
  Op2 b a
  -> NonEmpty a
  -> NonEmpty a
sortBy1 f =
  NonEmpty.sortBy (\a1 a2 -> view ordering (view op2 f a1 a2))