module Core.Data.List.NonEmpty
  ( replicate
  , (|>)
  , (|+)
  , (+|)
  , dropEnd
  , stripPrefix
  , zipPadS
  , zipPadM
  , zipPadLeftS
  , zipPadLeftM
  , zipPadWith
  ) where

import Data.Semigroup as S
import Data.Monoid as M
import qualified Data.List as List (replicate, stripPrefix)
import qualified Core.Data.List as List (
    (|>)
  , zipPadS
  , zipPadM
  , zipPadLeftS
  , zipPadLeftM
  , zipPadWith
  )
import Data.List.NonEmpty (NonEmpty (..))
import qualified Data.List.NonEmpty as NonEmpty
import Prelude hiding (replicate)

-- | Creates a list of length @n@, where every element is @x@.
replicate :: Int -> a -> NonEmpty a
replicate n x
  | n < 0 = error "replicate: tried to make NonEmpty with negative elements"
  | n == 0 = error "replicate: tried to make NonEmpty with 0 elements"
  | otherwise = NonEmpty.fromList $ List.replicate n x

-- | Appends an element to the list.
(|>) :: NonEmpty a -> a -> NonEmpty a
(x :| xs) |> y = x :| (xs List.|> y)

-- | Appends a regular list and a non-empty list.
(|+) :: [a] -> NonEmpty a -> NonEmpty a
[] |+ ys = ys
(x : xs) |+ ys = x :| (xs ++ NonEmpty.toList ys)

-- | Appends a non-empty list and a regular list.
(+|) :: NonEmpty a -> [a] -> NonEmpty a
(x :| xs) +| ys = x :| (xs ++ ys)

-- | Removes @n@ elements from the end of the list.
dropEnd :: Int -> NonEmpty a -> [a]
dropEnd n = reverse . NonEmpty.drop n . NonEmpty.reverse

-- | If the second list starts with the first, returns the part after.
-- Otherwise returns 'Nothing'.
stripPrefix :: (Eq a) => NonEmpty a -> NonEmpty a -> Maybe [a]
(x :| xs) `stripPrefix` (y :| ys)
  | x /= y = Nothing
  | otherwise = xs `List.stripPrefix` ys

-- | Zips the lists by appending elements. Won't discard elements at the
-- end of the longer list (so the result is as long as the longer list).
zipPadS :: (Semigroup a) => NonEmpty a -> NonEmpty a -> NonEmpty a
(x :| xs) `zipPadS` (y :| ys) = (x S.<> y) :| (xs `List.zipPadS` ys)

-- | Zips the lists by appending elements. Won't discard elements at the
-- end of the longer list (so the result is as long as the longer list).
zipPadM :: (Monoid a) => NonEmpty a -> NonEmpty a -> NonEmpty a
(x :| xs) `zipPadM` (y :| ys) = (x M.<> y) :| (xs `List.zipPadM` ys)

-- | Zips the lists by \prepending\ elements. The result is as long as
-- the longer list.
zipPadLeftS :: (Semigroup a) => NonEmpty a -> NonEmpty a -> NonEmpty a
(x :| xs) `zipPadLeftS` (y :| ys) = (x S.<> y) :| (xs `List.zipPadLeftS` ys)

-- | Zips the lists by \prepending\ elements. The result is as long as
-- the longer list.
zipPadLeftM :: (Monoid a) => NonEmpty a -> NonEmpty a -> NonEmpty a
(x :| xs) `zipPadLeftM` (y :| ys) = (x M.<> y) :| (xs `List.zipPadLeftM` ys)

-- | Zips the lists, appending 'mempty's to the shorter list instead of
-- discarding elements from the longer list (so the result is as long as
-- the longer list).
zipPadWith :: (Monoid a, Monoid b)
           => (a -> b -> c)
           -> NonEmpty a
           -> NonEmpty b
           -> NonEmpty c
zipPadWith f (x :| xs) (y :| ys) = f x y :| List.zipPadWith f xs ys