module Data.NonEmptyPrivate where

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

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


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

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)



newtype Zip f a = Zip {unZip :: f a}

instance Functor f => Functor (Zip f) where
   fmap f (Zip xs) = Zip $ fmap f xs

instance (C.Zip f, C.Repeat f) => Applicative (Zip f) where
   pure a = Zip $ C.repeat a
   Zip f <*> Zip x = Zip $ C.zipWith ($) f x



{- |
Always returns a rectangular list
by clipping all dimensions to the shortest slice.
Be aware that @transpose [] == repeat []@.
-}
transposeClip ::
   (Traversable f, C.Zip g, C.Repeat g) =>
   f (g a) -> g (f a)
transposeClip =
   unZip . Trav.sequenceA . fmap Zip


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