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


module Control.Functor.Expansive
  (
  -- * Expand
    Expansive (..)
  , uniteDichotomy
  ) where


import           Control.Applicative               (ZipList (ZipList))
import           Control.Functor.Dichotomous       (Dichotomous (ymotohcid),
                                                    These (..))
import           Data.Foldable                     (toList)
import           Data.Functor.Contravariant        (Contravariant (contramap))
import           Data.Functor.Product              (Product (..))
import qualified Data.IntMap                       as IntMap
import           Data.Kind                         (Type)
import qualified Data.Map                          as Map
import           Data.Maybe                        (isJust)
import           Data.Proxy                        (Proxy (Proxy))
import           Data.Semigroup                    (Option (..))
import qualified Data.Sequence                     as Seq
import qualified Data.Sequence.Internal            as Seq
import qualified Data.Vector                       as V
import           Data.Vector.Fusion.Bundle.Monadic (Bundle (..))
import qualified Data.Vector.Fusion.Bundle.Monadic as Bundle
import qualified Data.Vector.Fusion.Bundle.Size    as Bundle
import           Data.Vector.Fusion.Stream.Monadic (Step (..), Stream (..))
import           Data.Vector.Generic               (stream, unstream)


uniteDichotomy
  :: (Functor f, Expansive f, Dichotomous g)
  => f l -> f r -> f (Maybe (g l r))
uniteDichotomy x y = ymotohcid . Just <$> unite x y


-- | Partial inverse of Compactable
--
-- prop> expand (unite x y) = uniteDichotomy x y
-- prop> unite = emapThese id
-- prop> map Just = expand
-- prop> (\x -> unite x x) = map (\x -> These x x)
-- prop> emapThese f a b = map f (unite a b)
-- prop> unite (f <$> x) (g <$> y) = bimap f g <$> unite x y
-- prop> expand (unite x y) = swap <$> unite y x
-- prop> emapThese f a b = f <$> unite a b
-- prop> unite empty = map That
-- prop> flip unite empty = map This
-- prop> unite mempty = map That
-- prop> flip unite mempty = map This
class Expansive (f :: Type -> Type) where
  {-# MINIMAL unite | emapThese #-}

  expand :: f a -> f (Maybe a)
  default expand :: Functor f => f a -> f (Maybe a)
  expand = fmap Just
  {-# INLINABLE expand #-}

  unite :: f l -> f r -> f (These l r)
  unite = emapThese id
  {-# INLINABLE unite #-}

  unfilter :: (Bool -> a) -> f a -> f a
  unfilter f = emapMaybe $ f . isJust
  {-# INLINABLE unfilter #-}

  emapMaybe :: (Maybe b -> a) -> f b -> f a
  default emapMaybe :: Functor f => (Maybe b -> a) -> f b -> f a
  emapMaybe f = fmap f . expand
  {-# INLINABLE emapMaybe #-}

  econtramapMaybe :: Contravariant f => (a -> Maybe b) -> f b -> f a
  econtramapMaybe f = contramap f . expand
  {-# INLINABLE econtramapMaybe #-}

  emapThese :: (These l r -> a) -> f l -> f r -> f a
  default emapThese :: Functor f => (These l r -> a) -> f l -> f r -> f a
  emapThese f a b = f <$> unite a b
  {-# INLINABLE emapThese #-}

  econtramapThese :: Contravariant f => (a -> These l r) -> f l -> f r -> f a
  econtramapThese f l r = contramap f $ unite l r
  {-# INLINABLE econtramapThese #-}

  eapplyMaybe :: Applicative f => f (Maybe a -> b) -> f a -> f b
  eapplyMaybe fa = (fa <*>) . expand
  {-# INLINABLE eapplyMaybe #-}

  eapplyThese :: Applicative f => f (These l r -> a) -> f l -> f r -> f a
  eapplyThese fa = fmap (fa <*>) . unite
  {-# INLINABLE eapplyThese #-}

  ebindMaybe :: Applicative f => (f (Maybe b) -> a) -> f b -> f a
  ebindMaybe f x = pure . f $ expand x
  {-# INLINABLE ebindMaybe #-}

  ebindThese :: Applicative f => (f (These l r) -> a) -> f l -> f r -> f a
  ebindThese f x y = pure . f $ unite x y
  {-# INLINABLE ebindThese #-}


instance Expansive Maybe where
  unite (Just x) (Just y) = Just $ These x y
  unite (Just x) _        = Just $ This x
  unite _ (Just y)        = Just $ That y
  unite _ _               = Nothing
  {-# INLINABLE unite #-}


instance Expansive [] where
  unite xs []         = This <$> xs
  unite [] ys         = That <$> ys
  unite (x:xs) (y:ys) = These x y : unite xs ys
  {-# INLINABLE unite #-}


instance Expansive ZipList where
  unite (ZipList xs) (ZipList ys) = ZipList $ unite xs ys
  {-# INLINABLE unite #-}


instance Expansive Proxy where
  unite _ _ = Proxy
  {-# INLINABLE unite #-}


instance Expansive Option where
  unite (Option a) (Option b) = Option $ unite a b
  {-# INLINABLE unite #-}


-- instance (Applicative f, Applicative g) => Expansive (FP.Product f g) where
-- instance (Applicative f, Applicative g) => Expansive (Compose f g) where

instance Expansive Seq.Seq where
  unite xs (Seq.Seq Seq.EmptyT) = fmap This xs
  unite (Seq.Seq Seq.EmptyT) ys = fmap That ys
  unite xs ys                   = Seq.fromList $ unite (toList xs) (toList ys)
  {-# INLINABLE unite #-}


instance Monad m => Expansive (Bundle m v) where
  emapThese f Bundle{sElems = sa, sSize = na} Bundle{sElems = sb, sSize = nb}
    = Bundle.fromStream (emapThese f sa sb) (Bundle.larger na nb)
  {-# INLINABLE emapThese #-}


instance Monad m => Expansive (Stream m) where
  emapThese  f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing, False)
    where
    step (sa, sb, Nothing, False) = do
      r <- stepa sa
      return $ case r of
        Yield x sa' -> Skip (sa', sb, Just x,  False)
        Skip    sa' -> Skip (sa', sb, Nothing, False)
        Done        -> Skip (sa,  sb, Nothing, True)

    step (sa, sb, av, adone) = do
      r <- stepb sb
      return $ case r of
        Yield y sb' -> Yield (f $ maybe (That y) (`These` y) av)
                             (sa, sb', Nothing, adone)
        Skip sb'    -> Skip (sa, sb', av, adone)
        Done -> case (av, adone) of
          (Just x, False) -> Yield (f $ This x) (sa, sb, Nothing, adone)
          (_, True)       -> Done
          _               -> Skip (sa, sb, Nothing, False)


instance Expansive V.Vector where
  emapThese = emapThese'
    where
    emapThese' :: (These a b -> c) -> V.Vector a -> V.Vector b -> V.Vector c
    emapThese' f x y = unstream $ emapThese f (stream x) (stream y)
  {-# INLINABLE emapThese #-}


instance Expansive IntMap.IntMap where
  unite m n = IntMap.unionWith merge (IntMap.map This m) (IntMap.map That n)
    where merge (This a) (That b) = These a b
          merge _ _               = error "kimpossible"
  {-# INLINE unite #-}


instance Ord k => Expansive (Map.Map k) where
  unite m n = Map.unionWith merge (Map.map This m) (Map.map That n)
    where merge (This a) (That b) = These a b
          merge _ _               = error "kimpossible"
  {-# INLINE unite #-}


instance (Functor f, Functor g, Expansive f, Expansive g)
  => Expansive (Product f g) where
  unite (Pair a b) (Pair c d) = Pair (unite a c) (unite b d)
  {-# INLINE unite #-}