{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE InstanceSigs #-}

module Data.Ord.Ordering where

import Control.Applicative
    ( Applicative(liftA2, pure, (<*>)),
      (<$>),
      Alternative(empty, (<|>)) )
import Control.Lens
    ( view,
      iso,
      isn't,
      prism',
      (#),
      over,
      _Wrapped,
      Getting,
      Iso,
      Lens',
      Prism',
      Rewrapped,
      Wrapped(..) )
import Control.Monad.IO.Class ( MonadIO(..) )
import Control.Monad.Reader
import Control.Monad.State
    ( State, evalState )
import Control.Monad.Trans ( MonadTrans(..) )
import Data.Bool(bool)
import Data.Functor.Identity ( Identity(..) )
import Data.Functor.Contravariant ( Contravariant(contramap) )
import Data.Functor.Contravariant.Divisible
    ( Divisible(..) )
import Data.Functor ( Functor(fmap) )
import Data.Functor.Apply ( Apply((<.>)) )
import Data.Functor.Alt ( Alt((<!>)) )
import Data.Functor.Bind ( Bind((>>-)) )
import Data.List.NonEmpty ( NonEmpty(..) )
import Data.Profunctor ( Choice(..), Profunctor(dimap) )

-- $setup
-- >>> import Control.Lens
-- >>> import Control.Monad.State
-- >>> import Data.Eq
-- >>> import Data.Functor.Contravariant
newtype OrderT b f a =
  OrderT (a -> a -> f b)

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

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

type Order b a =
  OrderT b Identity a

order ::
  Iso
    (Order b a)
    (Order b a)
    (a -> a -> b)
    (a -> a -> b)
order =
  iso
    (\o a1 a2 -> runIdentity (view _Wrapped o a1 a2))
    (\f -> OrderT (\a1 a2 -> Identity (f a1 a2)))

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

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

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

instance Contravariant (OrderT b f) where
  contramap f (OrderT g) =
    OrderT (\a1 a2 -> g (f a1) (f a2))

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

-- |
-- >>> view order lt 1 1 :: Ordering
-- LT
--
-- >>> view order lt 1 2 :: Ordering
-- LT
--
-- >>> view order lt 2 1 :: Ordering
-- LT
--
-- prop> \x y -> (view order lt x y :: Ordering) == LT
-- +++ OK, passed 100 tests.
lt ::
  (Applicative f, AsOrdering b) =>
  OrderT b f a
lt =
  OrderT (\_ _ -> pure (_LT # ()))

-- |
--
-- >>> isLT LT
-- True
--
-- >>> isLT GT
-- False
--
-- >>> isLT EQ
-- False
isLT ::
  AsOrdering a =>
  a
  -> Bool
isLT =
  not . isn't _LT

-- |
--
-- >>> ifLT 1 2 LT
-- 2
-- >>> ifLT 1 2 GT
-- 1
-- >>> ifLT 1 2 EQ
-- 1
ifLT ::
  AsOrdering x =>
  a
  -> a
  -> x
  -> a
ifLT a1 a2 =
  bool a1 a2 . isLT

-- |
-- >>> view order eq 1 1 :: Ordering
-- EQ
--
-- >>> view order eq 1 2 :: Ordering
-- EQ
--
-- >>> view order eq 2 1 :: Ordering
-- EQ
--
-- prop> \x y -> (view order eq x y :: Ordering) == EQ
-- +++ OK, passed 100 tests.
eq ::
  (Applicative f, AsOrdering b) =>
  OrderT b f a
eq =
  OrderT (\_ _ -> pure (_EQ # ()))

-- |
--
-- >>> isEQ LT
-- False
--
-- >>> isEQ GT
-- False
--
-- >>> isEQ EQ
-- True
isEQ ::
  AsOrdering a =>
  a
  -> Bool
isEQ =
  not . isn't _EQ

-- |
--
-- >>> ifEQ 1 2 LT
-- 1
-- >>> ifEQ 1 2 GT
-- 1
-- >>> ifEQ 1 2 EQ
-- 2
ifEQ ::
  AsOrdering x =>
  a
  -> a
  -> x
  -> a
ifEQ a1 a2 =
  bool a1 a2 . isEQ

-- |
-- >>> view order gt 1 1 :: Ordering
-- GT
--
-- >>> view order gt 1 2 :: Ordering
-- GT
--
-- >>> view order gt 2 1 :: Ordering
-- GT
--
-- prop> \x y -> (view order gt x y :: Ordering) == GT
-- +++ OK, passed 100 tests.
gt ::
  (Applicative f, AsOrdering b) =>
  OrderT b f a
gt =
  OrderT (\_ _ -> pure (_GT # ()))

-- |
--
-- >>> isGT LT
-- False
--
-- >>> isGT GT
-- True
--
-- >>> isGT EQ
-- False
isGT ::
  AsOrdering a =>
  a
  -> Bool
isGT =
  not . isn't _GT

-- |
--
-- >>> ifGT 1 2 LT
-- 1
-- >>> ifGT 1 2 GT
-- 2
-- >>> ifGT 1 2 EQ
-- 1
ifGT ::
  AsOrdering x =>
  a
  -> a
  -> x
  -> a
ifGT a1 a2 =
  bool a1 a2 . isGT

-- |
--
-- prop> \x y -> view order ordOrder x y == x `compare` y
-- +++ OK, passed 100 tests.
ordOrder ::
  (Ord a, AsOrdering b, Applicative f) =>
  OrderT b f a
ordOrder =
  OrderT (\a1 a2 -> pure $
    case compare a1 a2 of
      LT ->
        _LT # ()
      GT ->
        _GT # ()
      EQ ->
        _EQ # ())

newtype MonadOrderT a f b =
  MonadOrderT (OrderT b f a)

monadOrder ::
  Iso
    (OrderT b f a)
    (OrderT b' f' a')
    (MonadOrderT a f b)
    (MonadOrderT a' f' b')
monadOrder =
  iso
    MonadOrderT
    (\(MonadOrderT x) -> x)

argument1 ::
  Applicative f =>
  MonadOrderT a f a
argument1 =
  MonadOrderT (OrderT (\a1 _ -> pure a1))

argument2 ::
  Applicative f =>
  MonadOrderT a f a
argument2 =
  MonadOrderT (OrderT (\_ a2 -> pure a2))

instance Functor f => Functor (MonadOrderT a f) where
  fmap :: Functor f => (a1 -> b) -> MonadOrderT a f a1 -> MonadOrderT a f b
  fmap f (MonadOrderT (OrderT g)) =
    MonadOrderT (OrderT (\a1 a2 -> fmap f (g a1 a2)))

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

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

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

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

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

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

instance MonadIO f => MonadIO (MonadOrderT a f) where
  liftIO a =
    MonadOrderT (OrderT (\ _ _ -> liftIO a))

instance MonadTrans (MonadOrderT a) where
  lift a =
    MonadOrderT (OrderT (\_ _ -> a))

newtype ProfunctorOrderT f a b =
  ProfunctorOrderT (OrderT b f a)

profunctorOrder ::
  Iso
    (OrderT b f a)
    (OrderT b' f' a')
    (ProfunctorOrderT f a b)
    (ProfunctorOrderT f' a' b')
profunctorOrder =
  iso
    ProfunctorOrderT
    (\(ProfunctorOrderT x) -> x)

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

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

instance Functor f => Functor (ProfunctorOrderT f a) where
  fmap f (ProfunctorOrderT (OrderT g)) =
    ProfunctorOrderT (OrderT (\a1 a2 -> fmap f (g a1 a2)))

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

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

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

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

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

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

instance MonadIO f => MonadIO (ProfunctorOrderT f a) where
  liftIO a =
    ProfunctorOrderT (OrderT (\ _ _ -> liftIO a))

appendOrder ::
  (Applicative f, Semigroup x) =>
  OrderT x f x
appendOrder =
  OrderT (\a1 a2 -> pure (a1 <> a2))

listOrder ::
  (Applicative f, AsOrdering b, Semigroup b) =>
  OrderT b f a
  -> OrderT b f [a]
listOrder (OrderT f) =
  let lists [] [] =
        pure (_EQ # ())
      lists [] (_:_) =
        pure (_LT # ())
      lists (_:_) [] =
        pure (_GT # ())
      lists (h1:t1) (h2:t2) =
        liftA2 (<>) (f h1 h2) (lists t1 t2)
  in  OrderT lists

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

bothOrder' ::
  Semigroup b =>
  (a -> b)
  -> Order b a
bothOrder' f =
  bothOrder (Identity . f)

-- |
--
-- >>> getPredicate (orderL _1 (Predicate even)) (1, "a")
-- False
--
-- >>> getPredicate (orderL _1 (Predicate even)) (2, "a")
-- True
--
-- >>> view order (orderL _1 ordOrder) (1, "a") (2, "b") :: Ordering
-- LT
--
-- >>> view order (orderL _1 ordOrder) (2, "a") (1, "b") :: Ordering
-- GT
--
-- >>> view order (orderL _1 ordOrder) (1, "a") (1, "b") :: Ordering
-- EQ
orderL ::
  Contravariant f =>
  Getting a s a
  -> f a
  -> f s
orderL =
  contramap . view

-- |
--
-- >>> view order (ordOrderL _1) (1, "a") (2, "b") :: Ordering
-- LT
--
-- >>> view order (ordOrderL _1) (2, "a") (1, "b") :: Ordering
-- GT
--
-- >>> view order (ordOrderL _1) (1, "a") (1, "b") :: Ordering
-- EQ
ordOrderL ::
  (Ord a, AsOrdering b) =>
  Getting a s a
  -> Order b s
ordOrderL x =
  orderL x ordOrder

-- |
--
-- >>> getPredicate (orderS (state (\s -> (1, reverse s))) (Predicate even)) "abc"
-- False
--
-- >>> getPredicate (orderS (state (\s -> (2, reverse s))) (Predicate even)) "abc"
-- True
--
-- >>> view order (orderS (state (\s -> (s + 1, s * 2))) ordOrder) 5 6 :: Ordering
-- LT
--
-- >>> view order (orderS (state (\s -> (s + 1, s * 2))) ordOrder) 5 4 :: Ordering
-- GT
--
-- >>> view order (orderS (state (\s -> (s + 1, s * 2))) ordOrder) 5 5 :: Ordering
-- EQ
orderS ::
  Contravariant f =>
  State a b
  -> f b
  -> f a
orderS =
  contramap . evalState

-- |
--
-- >>> view order (ordOrderS (state (\s -> (s + 1, s * 2)))) 5 6 :: Ordering
-- LT
--
-- >>> view order (ordOrderS (state (\s -> (s + 1, s * 2)))) 5 4 :: Ordering
-- GT
--
-- >>> view order (ordOrderS (state (\s -> (s + 1, s * 2)))) 5 5 :: Ordering
-- EQ
ordOrderS ::
  (Ord x, AsOrdering b) =>
  State a x
  -> Order b a
ordOrderS x =
  orderS x ordOrder

-- |
--
-- >>> getPredicate (orderR (reader (\r -> r + 1)) (Predicate even)) 1
-- True
--
-- >>> getPredicate (orderR (reader (\r -> r + 1)) (Predicate even)) 2
-- False
--
-- >>> view order (orderR (reader (\r -> r + 1)) ordOrder) 1 0 :: Ordering
-- GT
--
-- >>> view order (orderR (reader (\r -> r + 1)) ordOrder) 1 2 :: Ordering
-- LT
--
-- >>> view order (orderR (reader (\r -> r + 1)) ordOrder) 2 1 :: Ordering
-- GT
orderR ::
  Contravariant f =>
  Reader a b
  -> f b
  -> f a
orderR =
  contramap . runReader

-- |
--
-- >>> view order (ordOrderR (reader (\r -> r + 1))) 1 0 :: Ordering
-- GT
--
-- >>> view order (ordOrderR (reader (\r -> r + 1))) 1 2 :: Ordering
-- LT
--
-- >>> view order (ordOrderR (reader (\r -> r + 1))) 2 1 :: Ordering
-- GT
ordOrderR ::
  (Ord x, AsOrdering b, Applicative f) =>
  Reader a x
  -> OrderT b f a
ordOrderR x =
  orderR x ordOrder

-- |
--
-- >>> perRest (OrderT (\a1 a2 -> if a1 == 5 then Nothing else Just (a1 `compare` a2))) [5,1,2,3,5,6]
-- Nothing
--
-- >>> perRest (OrderT (\a1 a2 -> if a1 == 5 then Nothing else Just (a1 `compare` a2))) [5,1,2,3,6]
-- Just [(5,LT),(1,GT),(2,GT),(3,GT),(6,EQ)]
--
-- >>> perRest (OrderT (\a1 a2 -> if a1 == 0 then Nothing else Just (a1 `compare` a2))) [5,1,2,3,6]
-- Just [(5,LT),(1,GT),(2,GT),(3,GT),(6,EQ)]
--
-- >>> perRest (OrderT (\a1 a2 -> if a1 == 0 then Nothing else Just (a1 `compare` a2))) [4,5,1,2,3,6]
-- Just [(4,GT),(5,LT),(1,GT),(2,GT),(3,GT),(6,EQ)]
perRest ::
  (Applicative f, Monoid x) =>
  OrderT x f a
  -> [a]
  -> f [(a, x)]
perRest _ [] =
  pure []
perRest ord (h:r) =
  (\x y -> (h, x) : y) <$> foldr (\a b -> (<>) <$> view _Wrapped ord a h <*> b) (pure mempty) r <*> perRest ord r

-- |
--
-- >>> perRest' ordOrder [1,2,3,1,3,2,4] :: [(Int, Ordering)]
-- [(1,GT),(2,GT),(3,LT),(1,GT),(3,LT),(2,GT),(4,EQ)]
perRest' ::
  Monoid x =>
  Order x a
  -> [a]
  -> [(a, x)]
perRest' ord =
  runIdentity . perRest ord

-- | Returns a list of all elements in a list with at least one duplicate (equal according
-- to `Order`) in the remainder of the list.
--
-- >>> runIdentity (duplicates (ordOrder :: Order Ordering Int) [])
-- []
--
-- >>> runIdentity (duplicates (ordOrder :: Order Ordering Int) [1..10])
-- []
--
-- >>> runIdentity (duplicates (ordOrder :: Order Ordering Int) [1,2,3,1])
-- [(1,1 :| [])]
--
-- >>> runIdentity (duplicates (ordOrder :: Order Ordering Int) [1,2,3,1,4,5,1])
-- [(1,1 :| [1]),(1,1 :| [])]
--
-- >>> runIdentity (duplicates (ordOrder :: Order Ordering Int) [1,2,3,1,4,5,1,2,6,7,2,1])
-- [(1,1 :| [1,1]),(2,2 :| [2]),(1,1 :| [1]),(1,1 :| []),(2,2 :| [])]
duplicates ::
  (Monad f, AsOrdering b) =>
  OrderT b f a
  -> [a]
  -> f [(a, NonEmpty a)]
duplicates ord ts =
  let dups (_, []) = []
      dups (t, z:zs) = [(t, z:|zs)]
  in  fmap (>>= dups) (perRest (over monadOrder (\m -> (\a1 -> ifEQ [] [a1]) <$> argument1 <*> m) ord) ts)

-- | Returns a list of all elements in a list with at least one duplicate (equal according
-- to `Order`) in the remainder of the list.
--
-- >>> duplicates' (ordOrder :: Order Ordering Int) []
-- []
--
-- >>> duplicates' (ordOrder :: Order Ordering Int) [1..10]
-- []
--
-- >>> duplicates' (ordOrder :: Order Ordering Int) [1,2,3,1]
-- [(1,1 :| [])]
--
-- >>> duplicates' (ordOrder :: Order Ordering Int) [1,2,3,1,4,5,1]
-- [(1,1 :| [1]),(1,1 :| [])]
--
-- >>> duplicates' (ordOrder :: Order Ordering Int) [1,2,3,1,4,5,1,2,6,7,2,1]
-- [(1,1 :| [1,1]),(2,2 :| [2]),(1,1 :| [1]),(1,1 :| []),(2,2 :| [])]
duplicates' ::
  AsOrdering b =>
  Order b a
  -> [a]
  -> [(a, NonEmpty a)]
duplicates' ord =
  runIdentity . duplicates ord

-- | Asserts that the two given values (by the Eq instance) are equal in the `Order`
-- regardless of the function of the `Order`
--
-- >>> view order (areEqual 1 2 ordOrder) 3 4 :: Ordering
-- LT
--
-- >>> view order (areEqual 1 2 ordOrder) 4 3 :: Ordering
-- GT
--
-- >>> view order (areEqual 1 2 ordOrder) 3 3 :: Ordering
-- EQ
--
-- >>> view order (areEqual 1 2 ordOrder) 1 3 :: Ordering
-- LT
--
-- >>> view order (areEqual 1 2 ordOrder) 2 3 :: Ordering
-- LT
--
-- >>> view order (areEqual 1 2 ordOrder) 1 2 :: Ordering
-- EQ
--
-- >>> view order (areEqual 1 2 ordOrder) 2 1 :: Ordering
-- EQ
areEqual ::
  (Functor f, Eq a, AsOrdering b) =>
  a
  -> a
  -> OrderT b f a
  -> OrderT b f a
areEqual a1' a2' (OrderT f) =
  OrderT (\a1 a2 -> fmap (\e -> bool e (_EQ # ()) ((a1 == a1' && a2 == a2') || (a1 == a2' && a2 == a1'))) (f a1 a2))

-- | An alias for `areEqual`.
(.===.) ::
  (Applicative f, Ord a, AsOrdering b) =>
  a
  -> a
  -> OrderT b f a
a1 .===. a2 =
  areEqual a1 a2 ordOrder