module ListHelp 
    ( module Data.List
    , initOrNull
    , insertOver
    , swapIndices
    , (!?)
    , safeHead
    , safeUncons
    , errorHead
    , foldPairs
    , loopPairs
    , slideWindow
    
    , merge
    , mergeBy
    , mergeOn

    , takeWhileArb
    , takeWhileArb'
    , foldrWhileArb

    , wordsBy

    , findWithIx
    ) where
import Data.List
import Data.Ord
import FoldableHelp

loopPairs :: [a] -> [(a,a)]
loopPairs (a:as) = snd $ foldr f (Just a,[]) (a:as)
  where
    f x (Nothing,ps) = (Just x,ps)
    f x (Just y,ps)  = (Just x,(x,y):ps)
loopPairs [] = error "tried to make loop with empty list of elements"

foldPairs :: [a] -> [(a,a)]
foldPairs = snd . foldr f (Nothing, [])
  where
    f x (Nothing,ps) = (Just x,ps)
    f x (Just y,ps)  = (Just x,(x,y):ps)

initOrNull :: [a] -> [a]
initOrNull [] = []
initOrNull xs = init xs

insertOver :: Int -> a -> [a] -> [a]
insertOver i x xs = take i xs ++ x : drop (i+1) xs

swapIndices :: Int -> Int -> [a] -> [a]
swapIndices i j xs = insertOver i (xs !! j) . insertOver j (xs !! i) $ xs

--copied from Data.List.Extra:
-- | A total variant of the list index function `(!!)`.
--
-- > [2,3,4] !? 1    == Just 3
-- > [2,3,4] !? (-1) == Nothing
-- > []      !? 0    == Nothing
(!?) :: [a] -> Int -> Maybe a
xs !? n
  | n < 0     = Nothing
             -- Definition adapted from GHC.List
  | otherwise = foldr (\x r k -> case k of
                                   0 -> Just x
                                   _ -> r (k-1)) (const Nothing) xs n
{-# INLINABLE (!?) #-}

safeHead :: [a] -> Maybe a
safeHead (x:_) = Just x
safeHead _ = Nothing

safeUncons :: [a] -> Maybe (a,[a])
safeUncons (x:xs) = Just (x,xs)
safeUncons _ = Nothing

errorHead :: String -> [a] -> a
errorHead _ (x:_) = x
errorHead s [] = error s

slideWindow :: Int -> [a] -> [[a]]
slideWindow n = foldr (zipWith (:)) (repeat []) . take n . tails

merge :: Ord a => [a] -> [a] -> [a]
merge = mergeBy compare

mergeBy :: (a -> a -> Ordering) -> [a] -> [a] -> [a]
mergeBy _ xs [] = xs
mergeBy _ [] ys = ys
mergeBy f (x:xs) (y:ys)
    | f x y /= GT = x : mergeBy f xs (y:ys)
    | otherwise   = y : mergeBy f (x:xs) ys

mergeOn :: Ord b => (a -> b) -> [a] -> [a] -> [a]
mergeOn f xs = map snd . mergeBy (comparing fst) (g xs) . g
  where
    g = map (\x -> let y = f x in y `seq` (y,x))

-- | A generalised version of takeWhile
takeWhileArb :: (a -> Bool) -> (a -> b -> a) -> a -> [b] -> [b]
takeWhileArb _ _ _ [] = []
takeWhileArb t f y (x:xs) = case f y x of
    y' | t y' -> x : takeWhileArb t f y' xs
       | otherwise -> []

-- this may fuse better
-- should do tests
takeWhileArb' :: (a -> Bool) -> (a -> b -> a) -> a -> [b] -> [b]
takeWhileArb' t f a0 = snd . ssfold (t . fst) (\(a0',xs) x -> (f a0' x,x:xs)) (a0,[])

foldrWhileArb :: (b -> c -> c) -> c -> (a -> Bool) -> (a -> b -> a) -> a -> [b] -> c
foldrWhileArb _ z _ _ _ [] = z
foldrWhileArb h z t f y (x:xs) = case f y x of
    y' | t y' -> h x $ foldrWhileArb h z t f y' xs
       | otherwise -> h x z

wordsBy :: Eq a => a -> [a] -> [[a]]
wordsBy x xs = case dropWhile (==x) xs of
    [] -> []
    xs' -> w : wordsBy x xs''
      where
        (w,xs'') = break (==x) xs'

findWithIx :: (a -> Bool) -> [a] -> Maybe (Int,a)
findWithIx t = find (t . snd) . zip [0..]