module Data.NonEmptyPrivate where

import qualified Data.NonEmpty.Class as C
import qualified Data.Empty as Empty

import qualified Data.Sequence as Seq
import Data.Sequence (Seq, )

import qualified Data.Traversable as Trav
import qualified Data.Foldable as Fold
import qualified Data.List.HT as ListHT
import qualified Data.List as List
import Data.Traversable (Traversable, mapAccumL, mapAccumR)
import Data.Foldable (Foldable, )
import Control.Monad (Monad, return, (=<<), )
import Control.Applicative (Applicative, liftA2, pure, (<*>), )

import Data.Functor (Functor, fmap, )
import Data.Function (flip, const, ($), (.), )
import Data.Maybe (Maybe(Just, Nothing), maybe, mapMaybe, )
import Data.Ord (Ord, Ordering(GT), (<), (>), compare, comparing, )
import Data.Eq ((==), )
import Data.Tuple.HT (mapSnd, )
import Data.Tuple (fst, snd, )
import qualified Prelude as P
import Prelude (Eq, Show, Num, uncurry, )

import qualified Test.QuickCheck as QC


{-
We could also have (:!) as constructor,
but in order to import it unqualified we have to import 'T' unqualified, too,
and this would cause name clashes with locally defined types with name @T@.
-}
{- |
The type 'T' can be used for many kinds of list-like structures
with restrictions on the size.

* @T [] a@ is a lazy list containing at least one element.

* @T (T []) a@ is a lazy list containing at least two elements.

* @T Vector a@ is a vector with at least one element.
  You may also use unboxed vectors but the first element will be stored in a box
  and you will not be able to use many functions from this module.

* @T Maybe a@ is a list that contains one or two elements.

* @Maybe@ is isomorphic to @Optional Empty@.

* @T Empty a@ is a list that contains exactly one element.

* @T (T Empty) a@ is a list that contains exactly two elements.

* @Optional (T Empty) a@ is a list that contains zero or two elements.

* You can create a list type for every finite set of allowed list length
  by nesting Optional and NonEmpty constructors.
  If list length @n@ is allowed, then place @Optional@ at depth @n@,
  if it is disallowed then place @NonEmpty@.
  The maximm length is marked by @Empty@.
-}
data T f a = Cons { head :: a, tail :: f a }
   deriving (Eq, Ord)

instance (C.Show f, Show a) => Show (T f a) where
   showsPrec = C.showsPrec

instance (C.Show f) => C.Show (T f) where
   showsPrec p (Cons x xs) =
      P.showParen (p>5) $
      P.showsPrec 6 x . P.showString "!:" . C.showsPrec 5 xs


infixr 5 !:, `append`, `appendRight`, `appendLeft`

(!:) :: a -> f a -> T f a
(!:) = Cons


{- |
Force immediate generation of Cons.
-}
force :: T f a -> T f a
force x = Cons (head x) (tail x)


instance Functor f => Functor (T f) where
   fmap f (Cons x xs) = f x !: fmap f xs

instance Foldable f => Foldable (T f) where
   foldr f y (Cons x xs) = f x $ Fold.foldr f y xs
   foldl1 = foldl1
   foldr1 f (Cons x xs) =
      maybe x (f x) $
      Fold.foldr (\y -> Just . maybe y (f y)) Nothing xs
{-
   foldr1 f (Cons x xs) =
      case xs of
         [] -> x
         y:ys -> f x $ Fold.foldr1 f (Cons y ys)
-}


instance Traversable f => Traversable (T f) where
   sequenceA (Cons x xs) = liftA2 Cons x $ Trav.sequenceA xs

instance
   (Applicative f, C.Empty f, C.Cons f, C.Append f) =>
      Applicative (T f) where
   pure = singleton
   (<*>) = apply

instance (Monad f, C.Empty f, C.Cons f, C.Append f) =>
      Monad (T f) where
   return = singleton
   (>>=) = bind

instance (QC.Arbitrary a, C.Arbitrary f) => QC.Arbitrary (T f a) where
   arbitrary = liftA2 Cons QC.arbitrary C.arbitrary
   shrink (Cons x xs) = fmap (\(y, Aux ys) -> Cons y ys) $ QC.shrink (x, Aux xs)

newtype Aux f a = Aux (f a)

instance (C.Arbitrary f, QC.Arbitrary a) => QC.Arbitrary (Aux f a) where
   arbitrary = fmap Aux C.arbitrary
   shrink (Aux x) = fmap Aux $ C.shrink x

{- |
Implementation of 'Applicative.<*>' without the 'C.Empty' constraint
that is needed for 'Applicative.pure'.
-}
apply ::
   (Applicative f, C.Cons f, C.Append f) =>
   T f (a -> b) -> T f a -> T f b
apply (Cons f fs) (Cons x xs) =
   Cons (f x) (fmap f xs `C.append` (fs <*> C.cons x xs))

{- |
Implementation of 'Monad.>>=' without the 'C.Empty' constraint
that is needed for 'Monad.return'.
-}
bind ::
   (Monad f, C.Cons f, C.Append f) =>
   T f a -> (a -> T f b) -> T f b
bind (Cons x xs) k =
   appendRight (k x) (flatten . k =<< xs)


toList :: Foldable f => T f a -> [a]
toList (Cons x xs) = x : Fold.toList xs

flatten :: C.Cons f => T f a -> f a
flatten (Cons x xs) = C.cons x xs

fetch :: C.View f => f a -> Maybe (T f a)
fetch = fmap (uncurry Cons) . C.viewL


instance C.Cons f => C.Cons (T f) where
   cons = cons

cons :: C.Cons f => a -> T f a -> T f a
cons x0 (Cons x1 xs) = x0 !: C.cons x1 xs

-- snoc :: T f a -> a -> T f a
snocExtend :: Traversable f => f a -> a -> T f a
snocExtend xs y0 =
   uncurry Cons $ mapAccumR (\y x -> (x,y)) y0 xs


instance C.Empty f => C.Singleton (T f) where
   singleton = singleton

singleton :: C.Empty f => a -> T f a
singleton x = x !: C.empty


{-
This implementation needs quadratic time
with respect to the number of 'Cons'.
Maybe a linear time solution can be achieved using a type function
that maps a container type to the type of the reversed container.
-}
reverse :: (Traversable f, C.Reverse f) => T f a -> T f a
reverse (Cons x xs) = snocExtend (C.reverse xs) x

instance (Traversable f, C.Reverse f) => C.Reverse (T f) where
   reverse = reverse


mapHead :: (a -> a) -> T f a -> T f a
mapHead f (Cons x xs) = f x !: xs

mapTail :: (f a -> g a) -> T f a -> T g a
mapTail f (Cons x xs) = x !: f xs

init :: (C.Zip f, C.Cons f) => T f a -> f a
init (Cons x xs) = C.zipWith const (C.cons x xs) xs

last :: (Foldable f) => T f a -> a
last = foldl1 (flip const)

foldl1 :: (Foldable f) => (a -> a -> a) -> T f a -> a
foldl1 f (Cons x xs) = Fold.foldl f x xs

{- |
It holds:

> foldl1Map g f = foldl1 f . fmap g

but 'foldl1Map' does not need a 'Functor' instance.
-}
foldl1Map :: (Foldable f) => (a -> b) -> (b -> b -> b) -> T f a -> b
foldl1Map g f (Cons x xs) = Fold.foldl (\b a -> f b (g a)) (g x) xs


-- | maximum is a total function
maximum :: (Ord a, Foldable f) => T f a -> a
maximum = foldl1 P.max

-- | minimum is a total function
minimum :: (Ord a, Foldable f) => T f a -> a
minimum = foldl1 P.min

-- | maximumBy is a total function
maximumBy :: (Foldable f) => (a -> a -> Ordering) -> T f a -> a
maximumBy f = foldl1 (\x y -> case f x y of P.LT -> y; _ -> x)

-- | minimumBy is a total function
minimumBy :: (Foldable f) => (a -> a -> Ordering) -> T f a -> a
minimumBy f = foldl1 (\x y -> case f x y of P.GT -> y; _ -> x)

-- | maximumKey is a total function
maximumKey :: (Ord b, Foldable f) => (a -> b) -> T f a -> a
maximumKey f =
   snd .
   foldl1Map (attachKey f)
      (\ky0 ky1 -> if fst ky0 < fst ky1 then ky1 else ky0)

-- | minimumKey is a total function
minimumKey :: (Ord b, Foldable f) => (a -> b) -> T f a -> a
minimumKey f =
   snd .
   foldl1Map (attachKey f)
      (\ky0 ky1 -> if fst ky0 > fst ky1 then ky1 else ky0)

-- | maximumKey is a total function
_maximumKey :: (Ord b, Foldable f, Functor f) => (a -> b) -> T f a -> a
_maximumKey f =
   snd . maximumBy (comparing fst) . fmap (attachKey f)

-- | minimumKey is a total function
_minimumKey :: (Ord b, Foldable f, Functor f) => (a -> b) -> T f a -> a
_minimumKey f =
   snd . minimumBy (comparing fst) . fmap (attachKey f)

attachKey :: (a -> b) -> a -> (b, a)
attachKey f a = (f a, a)

-- | sum does not need a zero for initialization
sum :: (Num a, Foldable f) => T f a -> a
sum = foldl1 (P.+)

-- | product does not need a one for initialization
product :: (Num a, Foldable f) => T f a -> a
product = foldl1 (P.*)


instance (C.Cons f, C.Append f) => C.Append (T f) where
   append = append

append :: (C.Cons f, C.Append f) => T f a -> T f a -> T f a
append xs ys = appendRight xs (flatten ys)

appendRight :: (C.Append f) => T f a -> f a -> T f a
appendRight (Cons x xs) ys = Cons x (C.append xs ys)

appendLeft ::
   (C.Append f, C.View f, C.Cons f) =>
   f a -> T f a -> T f a
appendLeft xt yt =
   force $
   case C.viewL xt of
      Nothing -> yt
      Just (x,xs) -> Cons x $ C.append xs $ flatten yt


{- |
generic variants:
'Data.Monoid.HT.cycle' or better @Semigroup.cycle@
-}
cycle :: (C.Cons f, C.Append f) => T f a -> T f a
cycle x =
   let y = append x y
   in  y


instance (C.Zip f) => C.Zip (T f) where
   zipWith = zipWith

zipWith :: (C.Zip f) => (a -> b -> c) -> T f a -> T f b -> T f c
zipWith f (Cons a as) (Cons b bs) = Cons (f a b) (C.zipWith f as bs)


instance (C.Repeat f) => C.Repeat (T f) where
   repeat a = Cons a $ C.repeat a


instance (C.Sort f, Insert f) => C.Sort (T f) where
   sortBy = sortBy

{- |
If you nest too many non-empty lists
then the efficient merge-sort (linear-logarithmic runtime)
will degenerate to an inefficient insert-sort (quadratic runtime).
-}
sortBy :: (C.Sort f, Insert f) => (a -> a -> Ordering) -> T f a -> T f a
sortBy f (Cons x xs) =
   insertBy f x $ C.sortBy f xs

sort :: (Ord a, C.Sort f, Insert f) => T f a -> T f a
sort = sortBy compare



class Insert f where
   insertBy :: (a -> a -> Ordering) -> a -> f a -> T f a

instance (Insert f) => Insert (T f) where
   insertBy f y xt@(Cons x xs) =
      uncurry Cons $
      case f y x of
         GT -> (x, insertBy f y xs)
         _ -> (y, xt)

instance Insert Empty.T where
   insertBy _ x Empty.Cons = Cons x Empty.Cons

instance Insert [] where
   insertBy f y xt =
      uncurry Cons $
      case xt of
         [] -> (y, xt)
         x:xs ->
            case f y x of
               GT -> (x, List.insertBy f y xs)
               _ -> (y, xt)

instance Insert Maybe where
   insertBy f y mx =
      uncurry Cons $
      case mx of
         Nothing -> (y, Nothing)
         Just x ->
            mapSnd Just $
            case f y x of
               GT -> (x, y)
               _ -> (y, x)

instance Insert Seq where
   {-
   If we assume a sorted list
   we could do binary search for the splitting point.
   -}
   insertBy f y xt =
      uncurry Cons $
      case Seq.spanl ((GT ==) . f y) xt of
         (ys,zs) ->
            case Seq.viewl ys of
               Seq.EmptyL -> (y, xt)
               w Seq.:< ws -> (w, ws Seq.>< y Seq.<| zs)


{- |
Insert an element into an ordered list while preserving the order.
The first element of the resulting list is returned individually.
We need this for construction of a non-empty list.
-}
insert :: (Ord a, Insert f, C.Sort f) => a -> f a -> T f a
insert = insertBy compare



class Functor f => RemoveEach f where
   removeEach :: T f a -> T f (a, f a)

instance RemoveEach [] where
   removeEach (Cons x xs) =
      Cons (x, xs) (fmap (mapSnd (x:)) $ ListHT.removeEach xs)

instance RemoveEach Empty.T where
   removeEach (Cons x Empty.Cons) = Cons (x, Empty.Cons) Empty.Cons

instance RemoveEach f => RemoveEach (T f) where
   removeEach (Cons x xs) =
      Cons (x, xs) (fmap (mapSnd (x !:)) $ removeEach xs)

instance RemoveEach Maybe where
   removeEach (Cons x0 xs) =
      (\ ~(a,b) -> Cons (x0, a) b) $
      case xs of
         Nothing -> (Nothing, Nothing)
         Just x1 -> (Just x1, Just (x1, Just x0))


{-
It is somehow better than the variant in NonEmpty.Mixed,
since it can be applied to nested NonEmptys.
-}
class Tails f where
   tails :: (C.Cons g, C.Empty g) => f a -> T f (g a)

instance Tails [] where
   tails = tailsDefault

instance Tails Empty.T where
   tails Empty.Cons = Cons C.empty Empty.Cons

instance Tails f => Tails (T f) where
   tails (Cons x xs) =
      case tails xs of
         xss -> Cons (C.cons x $ head xss) xss

instance Tails Maybe where
   tails xs =
      force $
      case xs of
         Nothing -> Cons C.empty Nothing
         Just x -> Cons (C.cons x C.empty) (Just C.empty)

instance Tails Seq where
   tails = tailsDefault

tailsDefault ::
   (C.Cons f, C.Empty f, C.View f, Tails f,
    C.Cons g, C.Empty g) =>
   f a -> T f (g a)
tailsDefault xt =
   force $
   case C.viewL xt of
      Nothing -> Cons C.empty C.empty
      Just (x, xs) ->
         case tails xs of
            xss -> cons (C.cons x $ head xss) xss


{-
Not exorted by NonEmpty.
I think the transposeClip function is better.
-}
class TransposeOuter f where
   transpose :: TransposeInner g => f (g a) -> g (f a)

instance TransposeOuter [] where
   transpose =
      let go [] = transposeStart
          go (xs : xss) = zipHeadTail xs $ go xss
      in  go

{-
We cannot define this instance,
because @transpose ([] !: [2] !: []) = [2 !: []]@

instance TransposeOuter f => TransposeOuter (T f) where
   transpose =
      let go (Cons xs xss) = zipHeadTail xs $ go xss
      in  go
-}

class TransposeInner g where
   transposeStart :: g a
   zipHeadTail :: (C.Singleton f, C.Cons f) => g a -> g (f a) -> g (f a)

instance TransposeInner [] where
   transposeStart = []
   zipHeadTail =
      let go (x:xs) (ys:yss) = C.cons x ys : go xs yss
          go [] yss = yss
          go xs [] = fmap C.singleton xs
      in  go

{-
We cannot define this instance,
because @transpose ([] :: [NonEmpty.T [] Int]) = []@,
but in order to satisfy the types it must be ([] !: []).

instance TransposeInner f => TransposeInner (T f) where
   transposeStart = Cons ??? transposeStart
   zipHeadTail (Cons x xs) (Cons ys yss) =
      Cons (C.cons x ys) (zipHeadTail xs yss)
-}

{-
transpose :: [[a]] -> [[a]]
transpose =
   let go [] = []
       go (xs : xss) = zipHeadTail xs $ go xss
   in  go

zipHeadTail :: [a] -> [[a]] -> [[a]]
zipHeadTail (x:xs) (ys:yss) = (x:ys) : zipHeadTail xs yss
zipHeadTail [] yss = yss
zipHeadTail xs [] = fmap (:[]) xs
-}

transposePrelude :: [[a]] -> [[a]]
transposePrelude =
   let go [] = []
       go ([] : xss) = go xss
       go ((x:xs) : xss) =
          case ListHT.unzip $ mapMaybe ListHT.viewL xss of
             (ys, yss) -> (x : ys) : go (xs : yss)
   in  go

propTranspose :: [[P.Int]] -> P.Bool
propTranspose xs =
   List.transpose xs P.== transpose xs

propTransposePrelude :: [[P.Int]] -> P.Bool
propTransposePrelude xs =
   List.transpose xs P.== transposePrelude xs



scanl :: Traversable f => (b -> a -> b) -> b -> f a -> T f b
scanl f b =
   Cons b . snd .
   mapAccumL (\b0 -> (\b1 -> (b1,b1)) . f b0) b

scanr :: Traversable f => (a -> b -> b) -> b -> f a -> T f b
scanr f b =
   uncurry Cons .
   mapAccumR (\b0 -> flip (,) b0 . flip f b0) b

mapAdjacent ::
   (Traversable f) => (a -> a -> b) -> T f a -> f b
mapAdjacent f (Cons x xs) =
   snd $ mapAccumL (\a0 a1 -> (a1, f a0 a1)) x xs