{-# LANGUAGE CPP                        #-}
{-# LANGUAGE ConstrainedClassMethods    #-}
{-# LANGUAGE DefaultSignatures          #-}
{-# LANGUAGE DeriveFoldable             #-}
{-# LANGUAGE DeriveFunctor              #-}
{-# LANGUAGE DeriveGeneric              #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE KindSignatures             #-}
{-# LANGUAGE LambdaCase                 #-}
{-# LANGUAGE StandaloneDeriving         #-}
{-# LANGUAGE TypeOperators              #-}
{-# LANGUAGE TypeSynonymInstances       #-}


module Control.Functor.Dichotomous
  (
  -- Dichotomy
    Dichotomous (..)
  , hushLeft, hushRight, swap
  -- AltSum
  , AltSum (..)
  , mfold', mlefts, mrights
  , flipThese
  -- All types
  , LeftOnly(..)
  , LeftOrBoth(..)
  , RightOnly(..)
  , RightOrBoth(..)
  , These(..)
  , None (..)
  , MaybeBoth(..)
  , MaybeRight(..)
  , MaybeRightOrBoth(..)
  , MaybeLeft(..)
  , MaybeLeftOrBoth(..)
  , MaybeEither(..)
  , TheseOrNot(..)
  ) where


import           Control.Applicative (Alternative (empty, (<|>)))
import           Data.Bifoldable     (Bifoldable (bifoldMap))
import           Data.Kind           (Type)

import           Data.These          (These (..))
import           GHC.Generics        (Generic)


{-| Dichotomous is about types that are injective to (Maybe (These a b))
   In other words a + b + (a * b) + 1. Therefore
         ab          (,)
     b               LeftOnly
     b   ab          LeftOrBoth
     a               RightOnly
     a   ab          RightOrBoth
     a b             Either
     a b ab          These
   1                 None
   1     ab          MaybeBoth
   1   b             MaybeRight
   1   b ab          MaybeRightOrBoth
   1 a               MaybeLeft
   1 a   ab          MaybeLeftOrBoth
   1 a b             MaybeEither
   1 a b ab          TheseOrNot
-}


class Dichotomous (f :: Type -> Type -> Type) where
  dichotomy :: f a b -> Maybe (These a b)
  ymotohcid :: Maybe (These a b) -> Maybe (f a b)


instance Dichotomous (,) where
  dichotomy (x,y) = Just $ These x y
  {-# INLINE dichotomy #-}
  ymotohcid = \case Just (These x y) -> Just (x, y); _ -> Nothing
  {-# INLINE ymotohcid #-}


newtype LeftOnly a b = LeftOnly { unLeftOnly :: a }
  deriving (Eq, Ord, Show, Read, Generic)


instance Dichotomous LeftOnly where
  dichotomy (LeftOnly l) = Just (This l)
  {-# INLINE dichotomy #-}
  ymotohcid = \case
    Just (This l)    -> Just (LeftOnly l)
    Just (These l _) -> Just (LeftOnly l)
    _                -> Nothing
  {-# INLINE ymotohcid #-}


data LeftOrBoth a b = Left' a | LBoth a b
  deriving (Eq, Ord, Show, Read, Generic)


instance Dichotomous LeftOrBoth where
  dichotomy (Left' l)   = Just (This l)
  dichotomy (LBoth l r) = Just (These l r)
  {-# INLINE dichotomy #-}
  ymotohcid = \case
    Just (This l)    -> Just (Left' l)
    Just (These l r) -> Just (LBoth l r)
    _                -> Nothing
  {-# INLINE ymotohcid #-}


newtype RightOnly a b = RightOnly { unRightOnly :: b }
  deriving (Eq, Ord, Show, Read, Generic)


instance Dichotomous RightOnly where
  dichotomy (RightOnly r) = Just (That r)
  {-# INLINE dichotomy #-}
  ymotohcid = \case
    Just (That r)    -> Just (RightOnly r)
    Just (These _ r) -> Just (RightOnly r)
    _                -> Nothing
  {-# INLINE ymotohcid #-}


data RightOrBoth a b = Right' b | RBoth a b
  deriving (Eq, Ord, Show, Read, Generic)


instance Dichotomous RightOrBoth where
  dichotomy (Right' r)  = Just (That r)
  dichotomy (RBoth l r) = Just (These l r)
  {-# INLINE dichotomy #-}
  ymotohcid = \case
    Just (That r)    -> Just (Right' r)
    Just (These l r) -> Just (RBoth l r)
    _                -> Nothing
  {-# INLINE ymotohcid #-}


instance Dichotomous Either where
  dichotomy = Just . \case Left l -> This l; Right r -> That r
  {-# INLINABLE dichotomy #-}
  ymotohcid = \case Just (This l) -> Just (Left l)
                    Just (That r) -> Just (Right r)
                    _             -> Nothing
  {-# INLINEABLE ymotohcid #-}


instance Dichotomous These where
  dichotomy = Just
  {-# INLINE dichotomy #-}
  ymotohcid = id
  {-# INLINE ymotohcid #-}


data None a b = None
  deriving (Eq, Ord, Show, Read, Generic)


instance Dichotomous None where
  dichotomy _ = Nothing
  {-# INLINE dichotomy #-}
  ymotohcid _ = Just None
  {-# INLINE ymotohcid #-}


newtype MaybeBoth a b = MaybeBoth { unMaybeOrBoth :: Maybe (a,b) }
  deriving (Eq, Ord, Show, Read, Generic)


instance Dichotomous MaybeBoth where
  dichotomy x = case unMaybeOrBoth x of
    Nothing    -> Nothing
    Just (a,b) -> Just (These a b)
  {-# INLINE dichotomy #-}
  ymotohcid = Just . MaybeBoth . \case
    Just (These a b) -> Just (a,b)
    _                -> Nothing
  {-# INLINE ymotohcid #-}


data MaybeRight a b = MRNothing | MRight b
  deriving (Eq, Ord, Show, Read, Generic)


instance Dichotomous MaybeRight where
  dichotomy = \case
    MRNothing -> Nothing
    MRight b  -> Just (That b)
  {-# INLINE dichotomy #-}
  ymotohcid = Just . \case
    Just (That b)    -> MRight b
    Just (These _ b) -> MRight b
    _                -> MRNothing
  {-# INLINE ymotohcid #-}


data MaybeRightOrBoth a b = MRBNothing | MRBRight b | MRBoth a b
  deriving (Eq, Ord, Show, Read, Generic)


instance Dichotomous MaybeRightOrBoth where
  dichotomy = \case
    MRBNothing -> Nothing
    MRBRight b -> Just (That b)
    MRBoth a b -> Just (These a b)
  {-# INLINE dichotomy #-}
  ymotohcid = Just . \case
    Just (That b)    -> MRBRight b
    Just (These a b) -> MRBoth a b
    _                -> MRBNothing
  {-# INLINE ymotohcid #-}


data MaybeLeft a b = MLNothing | MLeft a
  deriving (Eq, Ord, Show, Read, Generic)


instance Dichotomous MaybeLeft where
  dichotomy = \case
    MLNothing -> Nothing
    MLeft b   -> Just (This b)
  {-# INLINE dichotomy #-}
  ymotohcid = Just . \case
    Just (This a)    -> MLeft a
    Just (These a _) -> MLeft a
    _                -> MLNothing
  {-# INLINE ymotohcid #-}


data MaybeLeftOrBoth a b = MLBNothing | MLBLeft a | MLBoth a b
  deriving (Eq, Ord, Show, Read, Generic)


instance Dichotomous MaybeLeftOrBoth where
  dichotomy = \case
    MLBNothing -> Nothing
    MLBLeft a  -> Just (This a)
    MLBoth a b -> Just (These a b)
  {-# INLINE dichotomy #-}
  ymotohcid = Just . \case
    Just (This a)    -> MLBLeft a
    Just (These a b) -> MLBoth a b
    _                -> MLBNothing
  {-# INLINE ymotohcid #-}


data MaybeEither a b = MENothing | MELeft a | MERight b
  deriving (Eq, Ord, Show, Read, Generic)


instance Dichotomous MaybeEither where
  dichotomy = \case
    MELeft a  -> Just (This a)
    MERight b -> Just (That b)
    _         -> Nothing
  {-# INLINE dichotomy #-}
  ymotohcid = Just . \case
    Just (This a) -> MELeft a
    Just (That a) -> MERight a
    _             -> MENothing
  {-# INLINE ymotohcid #-}


data TheseOrNot a b = This' a | That' b | These' a b | Not
  deriving (Eq, Ord, Show, Read, Generic)


instance Dichotomous TheseOrNot where
  dichotomy = \case
    This' a    -> Just (This a)
    That' b    -> Just (That b)
    These' a b -> Just (These a b)
    Not        -> Nothing
  {-# INLINE dichotomy #-}
  ymotohcid = Just . \case
    Nothing          -> Not
    Just (This a)    -> This' a
    Just (That a)    -> That' a
    Just (These a b) -> These' a b
  {-# INLINE ymotohcid #-}


newtype AltSum f a = AltSum { unAltSum :: f a }
    deriving (Functor, Applicative, Alternative)
instance Alternative f => Semigroup (AltSum f a) where
    AltSum a <> AltSum b = AltSum (a <|> b)
instance Alternative f => Monoid (AltSum f a) where
    mempty = empty
    AltSum a `mappend` AltSum b = AltSum (a <|> b)


mfold' :: (Foldable f, Alternative m) => f a -> m a
mfold' = unAltSum . foldMap (AltSum . pure)


mlefts :: (Bifoldable f, Alternative m) => f a b -> m a
mlefts = unAltSum . bifoldMap (AltSum . pure) (const mempty)


mrights :: (Bifoldable f, Alternative m) => f a b -> m b
mrights = unAltSum . bifoldMap (const mempty) (AltSum . pure)


hushRight :: Dichotomous g => g l r -> Maybe l
hushRight d = case dichotomy d of
  Just (This x)    -> Just x
  Just (These x _) -> Just x
  _                -> Nothing


hushLeft :: Dichotomous g => g l r -> Maybe r
hushLeft d = case dichotomy d of
  Just (That x)    -> Just x
  Just (These _ x) -> Just x
  _                -> Nothing


flipThese :: These a b -> These b a
flipThese = \case This x -> That x; That x -> This x; These x y -> These y x
{-# INLINABLE flipThese #-}


swap :: Dichotomous g => g a b -> Maybe (g b a)
swap g = ymotohcid $ case dichotomy g of
  Nothing          -> Nothing
  Just (This a)    -> Just (That a)
  Just (That b)    -> Just (This b)
  Just (These a b) -> Just (These b a)
{-# INLINABLE swap #-}