{-# OPTIONS_GHC -Wno-unused-imports #-}
module IntMapHelp
    ( module Data.Strict.IntMap
    , newKey
    , insertNewKey
    , insertWithNewKeys
    , unsafeSwapKeys
    , safeSwapKeys
    , findIndex
    ) where
import LensHelp
--import Data.List hiding (foldr,insert)
import Data.Strict.IntMap
import Data.Strict.Containers.Lens
{- | Find a key value one higher than any key in the map, or zero if the map is
 - empty -}
newKey :: IntMap a -> Int
newKey = maybe 0 ((+ 1) . fst) . lookupMax
{- | Insert an element with some new key. -}
insertNewKey :: a -> IntMap a -> IntMap a
insertNewKey x m = case lookupMax m of
    Nothing     -> singleton 0 x
    Just (k,_)  -> insert (k+1) x m
{- | Insert a list of values with new keys -}
insertWithNewKeys :: [a] -> IntMap a -> IntMap a
insertWithNewKeys xs m = union m $ fromList $ zip [newKey m..] xs
{- | Swaps two keys. Unsafe: assumes both keys exist. -}
unsafeSwapKeys :: Int -> Int -> IntMap a -> IntMap a
unsafeSwapKeys i j m = insert i (m ! j) $ insert j (m ! i) m

safeSwapKeys :: Int -> Int -> IntMap a -> IntMap a
safeSwapKeys i j m = m
    & at i .~ (m ^? ix j)
    & at j .~ (m ^? ix i)

-- ideally this should short circuit, I am not sure whether it does or not
-- though
findIndex :: (a -> Bool) -> IntMap a -> Maybe Int
findIndex t = foldrWithKey f Nothing
  where
    f _ _ (Just k) = Just k
    f k x _
        | t x = Just k
        | otherwise = Nothing