{-# OPTIONS_GHC -Wno-unused-imports #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

{- |
WARNING: orphan instances concerning Aeson classes and Data.Strict.IntMap datatypes have been introduced.
The warnings have been disabled.
-}
module IntMapHelp
--    ( module Data.Strict.IntMap
    (
    module Data.IntMap.Strict,
    newKey,
    insertNewKey,
    insertWithNewKeys,
    unsafeSwapKeys,
    safeSwapKeys,
    findIndex,
    cycleGT,
    cycleLT,
    invertIntMap,
    invertIntMapUnique,
) where

import Control.Applicative
import Data.Aeson
import Data.IntMap.Strict
import qualified Data.Map.Strict as M
import GHC.Generics
import LensHelp

--import Data.Strict.IntMap
--import Data.Strict.Containers.Lens

--import qualified Data.Strict.IntMap as IM
--import qualified Data.Strict.IntMap.Autogen.Strict as IM
--deriving instance Generic (IntMap a)

--instance ToJSON1 IntMap where
--    liftToJSON t tol = liftToJSON to' tol' . toList
--      where
--        to'  = liftToJSON2     toJSON toJSONList t tol
--        tol' = liftToJSONList2 toJSON toJSONList t tol
--
--    liftToEncoding t tol = liftToEncoding to' tol' . toList
--      where
--        to'  = liftToEncoding2     toEncoding toEncodingList t tol
--        tol' = liftToEncodingList2 toEncoding toEncodingList t tol
--
--instance ToJSON a => ToJSON (IntMap a) where
--    toJSON = toJSON1
--    toEncoding = toEncoding1
--instance FromJSON a => FromJSON (IntMap a) where
--    parseJSON = fmap fromList . parseJSON

--instance ToJSON a => ToJSON (IntMap a) where
--    toEncoding = genericToEncoding defaultOptions

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

cycleGT :: Int -> IntMap a -> Maybe (Int, a)
cycleGT i m = lookupGT i m <|> lookupMin m

cycleLT :: Int -> IntMap a -> Maybe (Int, a)
cycleLT i m = lookupLT i m <|> lookupMax m

invertIntMap :: Ord a => IntMap a -> M.Map a [Int]
invertIntMap = foldrWithKey (\k x -> M.insertWith (++) x [k]) mempty

invertIntMapUnique :: Ord a => IntMap a -> M.Map a Int
invertIntMapUnique = foldlWithKey' (\m k x -> M.insert x k m) mempty