    {-# LANGUAGE ScopedTypeVariables #-}

    -- | Utility functions for lists.

    module Agda.Utils.List where

    import Control.Arrow (first)

    import Data.Functor ((<$>))
    import Data.Function
    import qualified Data.List as List
    import Data.Maybe
    import qualified Data.Map as Map
    import qualified Data.Set as Set

    import Text.Show.Functions ()

    import qualified Agda.Utils.Bag as Bag

    import Agda.Utils.Tuple

    -- | Append a single element at the end.
    --   Time: O(length); use only on small lists.
    snoc :: [a] -> a -> [a]
    snoc xs x = xs ++ [x]

    -- | Case distinction for lists, with list first.
    --   Cf. 'Agda.Utils.Null.ifNull'.
    caseList :: [a] -> b -> (a -> [a] -> b) -> b
    caseList xs n c = listCase n c xs

    -- | Case distinction for lists, with list first.
    --   Cf. 'Agda.Utils.Null.ifNull'.
    caseListM :: Monad m => m [a] -> m b -> (a -> [a] -> m b) -> m b
    caseListM mxs n c = listCase n c =<< mxs

    -- | Case distinction for lists, with list last.
    listCase :: b -> (a -> [a] -> b) -> [a] -> b
    listCase n c []     = n
    listCase n c (x:xs) = c x xs

    -- | Head function (safe).
    headMaybe :: [a] -> Maybe a
    headMaybe = listToMaybe

    -- | Head function (safe). Returns a default value on empty lists.
    --
    -- > headWithDefault 42 []      = 42
    -- > headWithDefault 42 [1,2,3] = 1
    headWithDefault :: a -> [a] -> a
    headWithDefault def = fromMaybe def . headMaybe

    -- | Tail function (safe).
    tailMaybe :: [a] -> Maybe [a]
    tailMaybe = fmap snd . uncons

    -- | Tail function (safe).  Returns a default list on empty lists.
    tailWithDefault :: [a] -> [a] -> [a]
    tailWithDefault def = fromMaybe def . tailMaybe

    -- | Last element (safe).
    lastMaybe :: [a] -> Maybe a
    lastMaybe [] = Nothing
    lastMaybe xs = Just $ last xs

    -- | Last two elements (safe).
    last2 :: [a] -> Maybe (a, a)
    last2 (x : y : xs) = Just $ loop x y xs
      where
      loop x y []     = (x, y)
      loop x y (z:xs) = loop y z xs
    last2 _ = Nothing

    -- | Drop from the end of a list.
    --   @dropEnd n = reverse . drop n . reverse@
    --   (Forces the whole list even for @n==0@.)
    dropEnd :: forall a. Int -> [a] -> [a]
    dropEnd n = snd . foldr f (n, [])
      where
      f :: a -> (Int, [a]) -> (Int, [a])
      f x (n, xs)
        | n <= 0    = (0, x:xs)
        | otherwise = (n-1, xs)

    -- | Opposite of cons @(:)@, safe.
    uncons :: [a] -> Maybe (a, [a])
    uncons []     = Nothing
    uncons (x:xs) = Just (x,xs)

    -- | Maybe cons.   @mcons ma as = maybeToList ma ++ as@
    mcons :: Maybe a -> [a] -> [a]
    mcons ma as = maybe as (:as) ma

    -- | 'init' and 'last' in one go, safe.
    initLast :: [a] -> Maybe ([a],a)
    initLast []     = Nothing
    initLast (a:as) = Just $ loop a as where
      loop a []      = ([], a)
      loop a (b : bs) = mapFst (a:) $ loop b bs

    -- | init, safe.
    initMaybe :: [a] -> Maybe [a]
    initMaybe = fmap fst . initLast

    -- | Lookup function (partially safe).
    (!!!) :: [a] -> Int -> Maybe a
    []       !!! _         = Nothing
    (x : _)  !!! 0         = Just x
    (_ : xs) !!! n         = xs !!! (n - 1)

    -- | Lookup function with default value for index out of range.
    --   The name is chosen akin to 'Data.List.genericIndex'.
    indexWithDefault :: a -> [a] -> Int -> a
    indexWithDefault a []       _ = a
    indexWithDefault a (x : _)  0 = x
    indexWithDefault a (_ : xs) n = indexWithDefault a xs (n - 1)

    -- | Find an element satisfying a predicate and return it with its index.
    --   TODO: more efficient implementation!?
    findWithIndex :: (a -> Bool) -> [a] -> Maybe (a, Int)
    findWithIndex p as = headMaybe $ filter (p . fst) $ zip as [0..]

    -- | downFrom n = [n-1,..1,0]
    downFrom :: Integral a => a -> [a]
    downFrom n | n <= 0     = []
               | otherwise = let n' = n-1 in n' : downFrom n'

    -- | Update the first element of a list, if it exists.
    updateHead :: (a -> a) -> [a] -> [a]
    updateHead f [] = []
    updateHead f (a : as) = f a : as

    -- | Update the last element of a list, if it exists.
    updateLast :: (a -> a) -> [a] -> [a]
    updateLast f [] = []
    updateLast f [a] = [f a]
    updateLast f (a : as@(_ : _)) = a : updateLast f as

    -- | Update nth element of a list, if it exists.
    --   Precondition: the index is >= 0.
    updateAt :: Int -> (a -> a) -> [a] -> [a]
    updateAt _ f [] = []
    updateAt 0 f (a : as) = f a : as
    updateAt n f (a : as) = a : updateAt (n-1) f as

    -- | A generalized version of @partition@.
    --   (Cf. @mapMaybe@ vs. @filter@).
    mapEither :: (a -> Either b c) -> [a] -> ([b],[c])
    {-# INLINE mapEither #-}
    mapEither f xs = foldr (deal f) ([],[]) xs

    deal :: (a -> Either b c) -> a -> ([b],[c]) -> ([b],[c])
    deal f a ~(bs,cs) = case f a of
      Left  b -> (b:bs, cs)
      Right c -> (bs, c:cs)

    -- | Split off the largest suffix whose elements satisfy a predicate.
    --
    --   @spanEnd p xs = (ys, zs)@
    --   where @xs = ys ++ zs@
    --   and @all p zs@
    --   and @maybe True (not . p) (lastMaybe yz)@.
    spanEnd :: forall a. (a -> Bool) -> [a] -> ([a], [a])
    spanEnd p = snd . foldr f (True, ([], []))
      where
      f :: a -> (Bool, ([a], [a])) -> (Bool, ([a], [a]))
      f x (b', (xs, ys)) = (b, if b then (xs, x:ys) else (x:xs, ys))
        where b = b' && p x

    -- | A generalized version of @takeWhile@.
    --   (Cf. @mapMaybe@ vs. @filter@).
    takeWhileJust :: (a -> Maybe b) -> [a] -> [b]
    takeWhileJust p = loop
      where
        loop (a : as) | Just b <- p a = b : loop as
        loop _ = []

    -- | A generalized version of @span@.
    spanJust :: (a -> Maybe b) -> [a] -> ([b], [a])
    spanJust p = loop
      where
        loop (a : as) | Just b <- p a = mapFst (b :) $ loop as
        loop as                       = ([], as)

    -- | Partition a list into 'Nothing's and 'Just's.
    --   @'mapMaybe' f = snd . partitionMaybe f@.
    partitionMaybe :: (a -> Maybe b) -> [a] -> ([a], [b])
    partitionMaybe f = loop
      where
        loop []       = ([], [])
        loop (a : as) = case f a of
          Nothing -> mapFst (a :) $ loop as
          Just b  -> mapSnd (b :) $ loop as

    -- | Like 'filter', but additionally return the last partition
    --   of the list where the predicate is @False@ everywhere.
    filterAndRest :: (a -> Bool) -> [a] -> ([a],[a])
    filterAndRest p = mapMaybeAndRest $ \ a -> if p a then Just a else Nothing

    -- | Like 'mapMaybe', but additionally return the last partition
    --   of the list where the function always returns @Nothing@.
    mapMaybeAndRest :: (a -> Maybe b) -> [a] -> ([b],[a])
    mapMaybeAndRest f = loop [] where
      loop acc = \case
        []                   -> ([], reverse acc)
        x:xs | Just y <- f x -> first (y:) $ loop [] xs
             | otherwise     -> loop (x:acc) xs

    -- | Drops from both lists simultaneously until one list is empty.
    dropCommon :: [a] -> [b] -> ([a],[b])
    dropCommon (x : xs) (y : ys) = dropCommon xs ys
    dropCommon xs ys = (xs, ys)

    -- | Sublist relation.
    isSublistOf :: Eq a => [a] -> [a] -> Bool
    isSublistOf []       ys = True
    isSublistOf (x : xs) ys =
      case dropWhile (x /=) ys of
        []     -> False
        (_:ys) -> isSublistOf xs ys

    type Prefix a = [a]
    type Suffix a = [a]

    -- | Check if a list has a given prefix. If so, return the list
    --   minus the prefix.
    stripPrefixBy :: (a -> a -> Bool) -> Prefix a -> [a] -> Maybe (Suffix a)
    stripPrefixBy eq = loop
      where
      loop []    rest = Just rest
      loop (_:_) []   = Nothing
      loop (p:pat) (r:rest)
        | eq p r    = loop pat rest
        | otherwise = Nothing

    -- | @stripSuffix suf xs = Just pre@ iff @xs = pre ++ suf@.
    stripSuffix :: Eq a => Suffix a -> [a] -> Maybe (Prefix a)
    stripSuffix [] = Just
    stripSuffix s  = stripReversedSuffix (reverse s)

    type ReversedSuffix a = [a]

    -- | @stripReversedSuffix rsuf xs = Just pre@ iff @xs = pre ++ reverse suf@.
    stripReversedSuffix :: forall a. Eq a => ReversedSuffix a -> [a] -> Maybe (Prefix a)
    stripReversedSuffix rs = final . foldr step (SSSStrip rs)
      where
      -- Step of the automaton (reading input from right to left).
      step :: a -> StrSufSt a -> StrSufSt a
      step x = \case
        SSSMismatch   -> SSSMismatch
        SSSResult xs  -> SSSResult (x:xs)
        SSSStrip []   -> SSSResult [x]
        SSSStrip (y:ys)
          | x == y    -> SSSStrip ys
          | otherwise -> SSSMismatch

      -- Output of the automaton.
      final :: StrSufSt a -> Maybe (Prefix a)
      final = \case
        SSSResult xs -> Just xs
        SSSStrip []  -> Just []
        _            -> Nothing  -- We have not stripped the whole suffix or encountered a mismatch.

    -- | Internal state for stripping suffix.
    data StrSufSt a
      = SSSMismatch                 -- ^ Error.
      | SSSStrip (ReversedSuffix a) -- ^ "Negative string" to remove from end. List may be empty.
      | SSSResult [a]               -- ^ "Positive string" (result). Non-empty list.

    -- | Split a list into sublists. Generalisation of the prelude function
    --   @words@.
    --
    --   > words xs == wordsBy isSpace xs
    wordsBy :: (a -> Bool) -> [a] -> [[a]]
    wordsBy p xs = yesP xs
        where
            yesP xs = noP (dropWhile p xs)

            noP []  = []
            noP xs  = ys : yesP zs
                where
                    (ys,zs) = break p xs

    -- | Chop up a list in chunks of a given length.
    chop :: Int -> [a] -> [[a]]
    chop _ [] = []
    chop n xs = ys : chop n zs
        where (ys,zs) = splitAt n xs

    -- | Chop a list at the positions when the predicate holds. Contrary to
    --   'wordsBy', consecutive separator elements will result in an empty segment
    --   in the result.
    --    > intercalate [x] (chopWhen (== x) xs) == xs
    chopWhen :: (a -> Bool) -> [a] -> [[a]]
    chopWhen p [] = []
    chopWhen p xs =
      case break p xs of
        (w, [])     -> [w]
        (w, [_])    -> [w, []]
        (w, _ : ys) -> w : chopWhen p ys

    -- | All ways of removing one element from a list.
    holes :: [a] -> [(a, [a])]
    holes []     = []
    holes (x:xs) = (x, xs) : map (id -*- (x:)) (holes xs)

    -- | Check whether a list is sorted.
    --
    -- Assumes that the 'Ord' instance implements a partial order.

    sorted :: Ord a => [a] -> Bool
    sorted [] = True
    sorted xs = and $ zipWith (<=) xs (tail xs)

    -- | Check whether all elements in a list are distinct from each
    -- other. Assumes that the 'Eq' instance stands for an equivalence
    -- relation.
    distinct :: Eq a => [a] -> Bool
    distinct []     = True
    distinct (x:xs) = x `notElem` xs && distinct xs

    -- | An optimised version of 'distinct'.
    --
    -- Precondition: The list's length must fit in an 'Int'.

    fastDistinct :: Ord a => [a] -> Bool
    fastDistinct xs = Set.size (Set.fromList xs) == length xs

    -- | Checks if all the elements in the list are equal. Assumes that
    -- the 'Eq' instance stands for an equivalence relation.
    allEqual :: Eq a => [a] -> Bool
    allEqual []       = True
    allEqual (x : xs) = all (== x) xs

    -- | Returns an (arbitrary) representative for each list element
    --   that occurs more than once.
    duplicates :: Ord a => [a] -> [a]
    duplicates = mapMaybe dup . Bag.groups . Bag.fromList
      where
        dup (a : _ : _) = Just a
        dup _           = Nothing

    -- | A variant of 'List.groupBy' which applies the predicate to consecutive
    -- pairs.

    groupBy' :: (a -> a -> Bool) -> [a] -> [[a]]
    groupBy' _ []           = []
    groupBy' p xxs@(x : xs) = grp x $ zipWith (\x y -> (p x y, y)) xxs xs
      where
      grp x ys = (x : map snd xs) : tail
        where (xs, rest) = span fst ys
              tail = case rest of
                       []            -> []
                       ((_, z) : zs) -> grp z zs

    -- | @'groupOn' f = 'groupBy' (('==') \`on\` f) '.' 'List.sortBy' ('compare' \`on\` f)@.

    groupOn :: Ord b => (a -> b) -> [a] -> [[a]]
    groupOn f = List.groupBy ((==) `on` f) . List.sortBy (compare `on` f)

    -- | @splitExactlyAt n xs = Just (ys, zs)@ iff @xs = ys ++ zs@
    --   and @genericLength ys = n@.
    splitExactlyAt :: Integral n => n -> [a] -> Maybe ([a], [a])
    splitExactlyAt 0 xs       = return ([], xs)
    splitExactlyAt n []       = Nothing
    splitExactlyAt n (x : xs) = mapFst (x :) <$> splitExactlyAt (n-1) xs

    -- | A generalised variant of 'elemIndex'.

    genericElemIndex :: (Eq a, Integral i) => a -> [a] -> Maybe i
    genericElemIndex x xs =
      listToMaybe $
      map fst $
      filter snd $
      zip [0..] $
      map (== x) xs

    -- | Requires both lists to have the same length.
    --
    --   Otherwise, @Nothing@ is returned.

    zipWith' :: (a -> b -> c) -> [a] -> [b] -> Maybe [c]
    zipWith' f = loop
      where
      loop []        []      = Just []
      loop (x : xs) (y : ys) = (f x y :) <$> loop xs ys
      loop []       (_ : _)  = Nothing
      loop (_ : _)  []       = Nothing

    -- | Like 'zipWith' but keep the rest of the second list as-is
    --   (in case the second list is longer).
    --
    -- @
    --   zipWithKeepRest f as bs == zipWith f as bs ++ drop (length as) bs
    -- @
    zipWithKeepRest :: (a -> b -> b) -> [a] -> [b] -> [b]
    zipWithKeepRest f = loop
      where
      loop []       bs       = bs
      loop as       []       = []
      loop (a : as) (b : bs) = f a b : loop as bs

    -- -- UNUSED; a better type would be
    -- -- zipWithTails :: (a -> b -> c) -> [a] -> [b] -> ([c], Either [a] [b])

    -- -- | Like zipWith, but returns the leftover elements of the input lists.
    -- zipWithTails :: (a -> b -> c) -> [a] -> [b] -> ([c], [a] , [b])
    -- zipWithTails f xs       []       = ([], xs, [])
    -- zipWithTails f []       ys       = ([], [] , ys)
    -- zipWithTails f (x : xs) (y : ys) = (f x y : zs , as , bs)
    --   where (zs , as , bs) = zipWithTails f xs ys

    -- | Efficient variant of 'nubBy' for finite lists.
    --
    -- Specification:
    --
    -- > nubOn f xs == 'nubBy' ((==) `'on'` f) xs.
    nubOn :: Ord b => (a -> b) -> [a] -> [a]
    nubOn tag =
      map snd
      . List.sortBy (compare `on` fst)
      . map (snd . head)
      . List.groupBy ((==) `on` fst)
      . List.sortBy (compare `on` fst)
      . map (\p@(_, x) -> (tag x, p))
      . zip [1..]

    -- | Efficient variant of 'nubBy' for finite lists.
    --
    -- Specification: For each list @xs@ there is a list @ys@ which is a
    -- permutation of @xs@ such that
    --
    -- > uniqOn f xs == 'nubBy' ((==) `'on'` f) ys.
    --
    -- Furthermore
    --
    -- > List.sortBy (compare `on` f) (uniqOn f xs) == uniqOn f xs.
    uniqOn :: Ord b => (a -> b) -> [a] -> [a]
    uniqOn key = Map.elems . Map.fromList . map (\ a -> (key a, a))

    -- | Compute the common suffix of two lists.
    commonSuffix :: Eq a => [a] -> [a] -> [a]
    commonSuffix xs ys = reverse $ (commonPrefix `on` reverse) xs ys

    -- | Compute the common prefix of two lists.
    commonPrefix :: Eq a => [a] -> [a] -> [a]
    commonPrefix [] _ = []
    commonPrefix _ [] = []
    commonPrefix (x:xs) (y:ys)
      | x == y    = x : commonPrefix xs ys
      | otherwise = []

    editDistanceSpec :: Eq a => [a] -> [a] -> Int
    editDistanceSpec [] ys = length ys
    editDistanceSpec xs [] = length xs
    editDistanceSpec (x : xs) (y : ys)
      | x == y    = editDistanceSpec xs ys
      | otherwise = 1 + minimum [ editDistanceSpec (x : xs) ys
                                , editDistanceSpec xs (y : ys)
                                , editDistanceSpec xs ys ]

    editDistance :: Eq a => [a] -> [a] -> Int
    editDistance xs ys = editD 0 0
      where xss = List.tails xs
            yss = List.tails ys
            tbl = Map.fromList [ ((i, j), editD' i j) | i <- [0..length xss - 1], j <- [0..length yss - 1] ]
            editD i j = tbl Map.! (i, j)
            editD' i j =
              case (xss !! i, yss !! j) of
                ([], ys) -> length ys
                (xs, []) -> length xs
                (x : xs, y : ys)
                  | x == y    -> editD (i + 1) (j + 1)
                  | otherwise -> 1 + minimum [ editD (i + 1) j, editD i (j + 1), editD (i + 1) (j + 1) ]