module Data.NonEmptyPrivate where

import qualified Data.NonEmpty.Class as C

import qualified Data.Traversable as Trav
import qualified Data.Foldable as Fold
import Data.Traversable (Traversable, )
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, )
import Data.Ord (Ord, Ordering(GT), compare, )
import Data.Tuple.HT (forcePair, )
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.

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

* @T (T Empty) a@ is a list that contains exactly two elements.
-}
data T f a = Cons { head :: a, tail :: f a }
   deriving (Eq, Ord, Show)


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

(!:) :: 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)


data Empty a = Empty
   deriving (Eq, Ord, Show)

instance Functor Empty where
   fmap _ Empty = Empty

instance Foldable Empty where
   foldr _ y Empty = y

instance Traversable Empty where
   sequenceA Empty = pure Empty

instance C.View Empty where
   viewL _ = Nothing

instance QC.Arbitrary (Empty a) where
   arbitrary = return Empty
   shrink _ = []


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


instance C.Empty Empty where
   empty = Empty

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

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

reverse :: (Foldable f, C.Cons f, C.Empty f) => T f a -> T f a
reverse (Cons x xs) =
   Fold.foldl (flip cons) (singleton x) xs

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


-- | 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

-- | 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)

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.Sort f) => C.Sort (T f) where
   sortBy = sortBy
   insertBy f y xt@(Cons x xs) =
      forcePair $
      case f y x of
         GT -> (x, uncurry Cons $ C.insertBy f y xs)
         _ -> (y, xt)

{- |
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) => (a -> a -> Ordering) -> T f a -> T f a
sortBy f (Cons x xs) =
   uncurry Cons $ C.insertBy f x $ C.sortBy f xs

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

insertBy ::
   (C.Sort f, C.Cons f) =>
   (a -> a -> Ordering) -> a -> T f a -> T f a
insertBy f y = uncurry cons . C.insertBy f y

insert ::
   (Ord a, C.Sort f, C.Cons f) =>
   a -> T f a -> T f a
insert = insertBy compare