{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RankNTypes #-}

module Waterdeep.Util
   ( deleteAt
   , on
   , nubOn
   , countOf
   , singular
   , mif
   , joinStrings
   , (<$$>)
   , whenM
   , unlessM
   , nextIntKey
   , intAtNext
   , walkIntMap
   ) where

import Control.Applicative
import Control.Monad.State
import Data.List
import Data.Function
import Data.Maybe
import Data.Monoid
import Lens.Family2 hiding ((&))
import Lens.Family2.State
import Lens.Family2.Stock

import qualified Data.IntMap as IM

deleteAt :: Int -> [a] -> [a]
deleteAt n l = left ++ drop 1 right
   where (left, right) = splitAt n l

nubOn :: Eq b => (a -> b) -> [a] -> [a]
nubOn f = nubBy ((==) `on` f)

countOf :: Num r => FoldLike (Sum r) a a' b b' -> (b -> Bool) -> a -> r
countOf l f = getSum . views l (\b -> if f b then Sum 1 else Sum 0)

-- |Turn a traversal into a lens. This is well-defined if and only if the traversal
-- always refers to exactly one element. With multiple elements, setting applies to
-- all and getting returns just the first value. With zero elements, setting is a
-- no-op and getting results in a runtime error.
singular :: Traversal a a' b b' -> Lens a a' b b'
singular t f b = (\a' -> b & t .~ a') `fmap` (f (unJust (firstOf t b)))
   where unJust (Just x) = x
         unJust Nothing  = error "singular: empty traversal"

mif :: Monoid m => Bool -> m -> m
mif c m = if c then m else mempty

joinStrings :: [String] -> String
joinStrings []      = "nothing"
joinStrings [x]     = x
joinStrings [x,y]   = x ++ " and " ++ y
joinStrings [x,y,z] = x ++ ", " ++ y ++ ", and " ++ z
joinStrings (x:xs)  = x ++ ", " ++ joinStrings xs

-- Like (<$>), but with the same fixity and precedence as ($)
(<$$>) :: Applicative f => (a -> b) -> f a -> f b
(<$$>) = (<$>)
infixr 0 <$$>

-- Like when/unless, except that the condition is an action yielding a Bool
-- rather than a pure Bool. Avoids the need for spurious condition variables.
whenM, unlessM :: Monad m => m Bool -> m () -> m ()
whenM   mc ma = do { c <- mc; when   c ma }
unlessM mc ma = do { c <- mc; unless c ma }

-- Returns the next available key in an IntMap, starting from 1 if the map is empty.
nextIntKey :: IM.IntMap a -> IM.Key
nextIntKey im = if IM.null im then 1 else 1 + fst (IM.findMax im)

-- This isn't really a lens since it doesn't refer to a constant location.
-- You don't get back what you put in, and setting twice has a different
-- effect than setting once. Nonetheless, I find it useful. It works like
-- intAt except that the (implied) key parameter always refers to the next
-- available key.
intAtNext :: Applicative f => (Maybe b -> f (Maybe b)) -> IM.IntMap b -> f (IM.IntMap b)
intAtNext f im = maybe im (flip (IM.insert (nextIntKey im)) im) <$> f Nothing

walkIntMap :: MonadState a m => Lens' a (IM.IntMap b)
           -> (IM.Key -> Lens' a (Maybe b) -> m r) -> m ()
walkIntMap l f = use l >>= \im -> unless (IM.null im) $ go (fst $ IM.findMin im) where
   go k = f k (l . intAt k) >> use l >>= maybe (return ()) (go . fst) . IM.lookupGT k