{-# LANGUAGE CPP #-}

-- |
-- Simple 'Data.Map'-based histogram.
-- A histogram counts occurrences of things, i.e. 'Histogram k' represents a mapping @k -> Int@.
-- Since it is backed by a 'Map' from 'Data.Map', it requires @k@ to have an @Ord@ instance.
module Data.Histogram
  ( Histogram,
    toMap,
    increment,
    decrement,
    Data.Histogram.lookup,
    (!),
    add,
    set,
    reset,
    zero,
    nonzero,
    size,
    empty,
    keys,
    mapKeys,
    singleton,
    singletonCount,
    split,
    splitLookup,
    isSubsetOf,
    isSubsetOfBy,
    disjoint,
    fromList,
    fromCountList,
    flatMap,
    toList,
    fromMap,
    unsafeFromMap,
  )
where

import qualified Data.Map.Strict as M
import Data.Maybe (fromMaybe)

-- | A simple 'Data.Map'-based histogram that counts occurrences of @k@.
newtype Histogram k = Histogram
  { -- | Convert to a histogram to a map of counts of all nonzero values
    Histogram k -> Map k Int
toMap :: M.Map k Int
  }
  deriving (Histogram k -> Histogram k -> Bool
(Histogram k -> Histogram k -> Bool)
-> (Histogram k -> Histogram k -> Bool) -> Eq (Histogram k)
forall k. Eq k => Histogram k -> Histogram k -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Histogram k -> Histogram k -> Bool
$c/= :: forall k. Eq k => Histogram k -> Histogram k -> Bool
== :: Histogram k -> Histogram k -> Bool
$c== :: forall k. Eq k => Histogram k -> Histogram k -> Bool
Eq, Int -> Histogram k -> ShowS
[Histogram k] -> ShowS
Histogram k -> String
(Int -> Histogram k -> ShowS)
-> (Histogram k -> String)
-> ([Histogram k] -> ShowS)
-> Show (Histogram k)
forall k. Show k => Int -> Histogram k -> ShowS
forall k. Show k => [Histogram k] -> ShowS
forall k. Show k => Histogram k -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Histogram k] -> ShowS
$cshowList :: forall k. Show k => [Histogram k] -> ShowS
show :: Histogram k -> String
$cshow :: forall k. Show k => Histogram k -> String
showsPrec :: Int -> Histogram k -> ShowS
$cshowsPrec :: forall k. Show k => Int -> Histogram k -> ShowS
Show)

instance Ord k => Semigroup (Histogram k) where
  Histogram Map k Int
m1 <> :: Histogram k -> Histogram k -> Histogram k
<> Histogram Map k Int
m2 = Map k Int -> Histogram k
forall k. Map k Int -> Histogram k
Histogram (Map k Int -> Histogram k) -> Map k Int -> Histogram k
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Int) -> Map k Int -> Map k Int -> Map k Int
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
M.unionWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) Map k Int
m1 Map k Int
m2

instance Ord k => Monoid (Histogram k) where
  mempty :: Histogram k
mempty = Map k Int -> Histogram k
forall k. Map k Int -> Histogram k
Histogram Map k Int
forall a. Monoid a => a
mempty
  mappend :: Histogram k -> Histogram k -> Histogram k
mappend = Histogram k -> Histogram k -> Histogram k
forall a. Semigroup a => a -> a -> a
(<>)

{-# INLINE clip #-}
clip :: Int -> Maybe Int
clip :: Int -> Maybe Int
clip Int
n
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 = Int -> Maybe Int
forall a. a -> Maybe a
Just Int
n
  | Bool
otherwise = Maybe Int
forall a. Maybe a
Nothing

-- | Construct a histogram from a map, removing all elements smaller than 1
fromMap :: M.Map k Int -> Histogram k
fromMap :: Map k Int -> Histogram k
fromMap = Map k Int -> Histogram k
forall k. Map k Int -> Histogram k
Histogram (Map k Int -> Histogram k)
-> (Map k Int -> Map k Int) -> Map k Int -> Histogram k
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Maybe Int) -> Map k Int -> Map k Int
forall a b k. (a -> Maybe b) -> Map k a -> Map k b
M.mapMaybe Int -> Maybe Int
clip

-- | Construct a histogram directly from a map, without checking if every element is above 1
unsafeFromMap :: M.Map k Int -> Histogram k
unsafeFromMap :: Map k Int -> Histogram k
unsafeFromMap = Map k Int -> Histogram k
forall k. Map k Int -> Histogram k
Histogram

-- | Increase a key's count by one
increment :: Ord k => k -> Histogram k -> Histogram k
increment :: k -> Histogram k -> Histogram k
increment k
k (Histogram Map k Int
m) = Map k Int -> Histogram k
forall k. Map k Int -> Histogram k
Histogram (Map k Int -> Histogram k) -> Map k Int -> Histogram k
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Int) -> k -> Int -> Map k Int -> Map k Int
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) k
k Int
1 Map k Int
m

-- | Decrease a key's count by one
decrement :: Ord k => k -> Histogram k -> Histogram k
decrement :: k -> Histogram k -> Histogram k
decrement k
k (Histogram Map k Int
m) = Map k Int -> Histogram k
forall k. Map k Int -> Histogram k
Histogram (Map k Int -> Histogram k) -> Map k Int -> Histogram k
forall a b. (a -> b) -> a -> b
$ (Int -> Maybe Int) -> k -> Map k Int -> Map k Int
forall k a. Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a
M.update Int -> Maybe Int
f' k
k Map k Int
m
  where
    f' :: Int -> Maybe Int
f' Int
n = Int -> Maybe Int
clip (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)

-- | Increase a key's count by an arbitrary number.
--   Can also be used to decrease by passing a negative value.
--   If the count falls below zero, it's set to 0.
add :: Ord k => Int -> k -> Histogram k -> Histogram k
add :: Int -> k -> Histogram k -> Histogram k
add Int
n k
k (Histogram Map k Int
m) = Map k Int -> Histogram k
forall k. Map k Int -> Histogram k
Histogram (Map k Int -> Histogram k) -> Map k Int -> Histogram k
forall a b. (a -> b) -> a -> b
$ (Maybe Int -> Maybe Int) -> k -> Map k Int -> Map k Int
forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
M.alter Maybe Int -> Maybe Int
f k
k Map k Int
m
  where
    f :: Maybe Int -> Maybe Int
f Maybe Int
nOld = Int -> Maybe Int
clip (Int -> Maybe Int) -> Int -> Maybe Int
forall a b. (a -> b) -> a -> b
$ Int -> Maybe Int -> Int
forall a. a -> Maybe a -> a
fromMaybe Int
0 Maybe Int
nOld Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n

-- | Set a key's count to an exact value.
--   Nonpositive numbers clip to 0.
set :: Ord k => Int -> k -> Histogram k -> Histogram k
set :: Int -> k -> Histogram k -> Histogram k
set Int
n k
k (Histogram Map k Int
m) = Map k Int -> Histogram k
forall k. Map k Int -> Histogram k
Histogram (Map k Int -> Histogram k) -> Map k Int -> Histogram k
forall a b. (a -> b) -> a -> b
$ (if Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 then (k -> Int -> Map k Int -> Map k Int)
-> Int -> k -> Map k Int -> Map k Int
forall a b c. (a -> b -> c) -> b -> a -> c
flip k -> Int -> Map k Int -> Map k Int
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Int
n else k -> Map k Int -> Map k Int
forall k a. Ord k => k -> Map k a -> Map k a
M.delete) k
k Map k Int
m

-- | Set a key's count to 0.
reset :: Ord k => k -> Histogram k -> Histogram k
reset :: k -> Histogram k -> Histogram k
reset k
k (Histogram Map k Int
m) = Map k Int -> Histogram k
forall k. Map k Int -> Histogram k
Histogram (Map k Int -> Histogram k) -> Map k Int -> Histogram k
forall a b. (a -> b) -> a -> b
$ k -> Map k Int -> Map k Int
forall k a. Ord k => k -> Map k a -> Map k a
M.delete k
k Map k Int
m

-- | Check whether a key has a count of at least 1.
nonzero :: Ord k => k -> Histogram k -> Bool
nonzero :: k -> Histogram k -> Bool
nonzero k
k (Histogram Map k Int
m) = k -> Map k Int -> Bool
forall k a. Ord k => k -> Map k a -> Bool
M.member k
k Map k Int
m

-- | Check whether a key has a count of 0
zero :: Ord k => k -> Histogram k -> Bool
zero :: k -> Histogram k -> Bool
zero k
k = Bool -> Bool
not (Bool -> Bool) -> (Histogram k -> Bool) -> Histogram k -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. k -> Histogram k -> Bool
forall k. Ord k => k -> Histogram k -> Bool
nonzero k
k

-- | Get the total number of elements in the map, i.e. the sum of all bins.
size :: Histogram k -> Int
size :: Histogram k -> Int
size = [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int) -> (Histogram k -> [Int]) -> Histogram k -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map k Int -> [Int]
forall k a. Map k a -> [a]
M.elems (Map k Int -> [Int])
-> (Histogram k -> Map k Int) -> Histogram k -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Histogram k -> Map k Int
forall k. Histogram k -> Map k Int
toMap

-- | Check whether a histogram is empty
empty :: Histogram k -> Bool
empty :: Histogram k -> Bool
empty = Map k Int -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null (Map k Int -> Bool)
-> (Histogram k -> Map k Int) -> Histogram k -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Histogram k -> Map k Int
forall k. Histogram k -> Map k Int
toMap

-- | Get a list of all non-zero keys.
keys :: Histogram k -> [k]
keys :: Histogram k -> [k]
keys = Map k Int -> [k]
forall k a. Map k a -> [k]
M.keys (Map k Int -> [k])
-> (Histogram k -> Map k Int) -> Histogram k -> [k]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Histogram k -> Map k Int
forall k. Histogram k -> Map k Int
toMap

-- | Applies a function to every key.
--   If two keys in the original map to the same value, their counts are combined.
mapKeys :: Ord k2 => (k1 -> k2) -> Histogram k1 -> Histogram k2
mapKeys :: (k1 -> k2) -> Histogram k1 -> Histogram k2
mapKeys k1 -> k2
f (Histogram Map k1 Int
m) = Map k2 Int -> Histogram k2
forall k. Map k Int -> Histogram k
Histogram (Map k2 Int -> Histogram k2) -> Map k2 Int -> Histogram k2
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Int) -> (k1 -> k2) -> Map k1 Int -> Map k2 Int
forall k2 a k1.
Ord k2 =>
(a -> a -> a) -> (k1 -> k2) -> Map k1 a -> Map k2 a
M.mapKeysWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) k1 -> k2
f Map k1 Int
m

-- | A histogram containing one key with a count of 1.
singleton :: k -> Histogram k
singleton :: k -> Histogram k
singleton k
k = Map k Int -> Histogram k
forall k. Map k Int -> Histogram k
Histogram (Map k Int -> Histogram k) -> Map k Int -> Histogram k
forall a b. (a -> b) -> a -> b
$ k -> Int -> Map k Int
forall k a. k -> a -> Map k a
M.singleton k
k Int
1

singletonCount :: Ord k => k -> Int -> Histogram k
singletonCount :: k -> Int -> Histogram k
singletonCount k
k Int
n
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 = Map k Int -> Histogram k
forall k. Map k Int -> Histogram k
Histogram (Map k Int -> Histogram k) -> Map k Int -> Histogram k
forall a b. (a -> b) -> a -> b
$ k -> Int -> Map k Int
forall k a. k -> a -> Map k a
M.singleton k
k Int
n
  | Bool
otherwise = Histogram k
forall a. Monoid a => a
mempty

-- | @isSubsetOfBy f h1 h2@ returns 'True' if every key in @h1@ compares to 'True' to its corresponding key in @h2@ by @f@.
isSubsetOfBy :: Ord k => (Int -> Int -> Bool) -> Histogram k -> Histogram k -> Bool
isSubsetOfBy :: (Int -> Int -> Bool) -> Histogram k -> Histogram k -> Bool
isSubsetOfBy Int -> Int -> Bool
f (Histogram Map k Int
h1) (Histogram Map k Int
h2) = (Int -> Int -> Bool) -> Map k Int -> Map k Int -> Bool
forall k a b.
Ord k =>
(a -> b -> Bool) -> Map k a -> Map k b -> Bool
M.isSubmapOfBy Int -> Int -> Bool
f Map k Int
h1 Map k Int
h2

-- | @isSubsetOf h1 h2@ returns 'True' if no key has a greater count in @h1@ than in @h2@.
isSubsetOf :: Ord k => Histogram k -> Histogram k -> Bool
isSubsetOf :: Histogram k -> Histogram k -> Bool
isSubsetOf = (Int -> Int -> Bool) -> Histogram k -> Histogram k -> Bool
forall k.
Ord k =>
(Int -> Int -> Bool) -> Histogram k -> Histogram k -> Bool
isSubsetOfBy Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(<=)

-- | Construct a histogram by counting occurrences in a list of keys.
fromList :: Ord k => [k] -> Histogram k
fromList :: [k] -> Histogram k
fromList = (k -> Histogram k -> Histogram k)
-> Histogram k -> [k] -> Histogram k
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr k -> Histogram k -> Histogram k
forall k. Ord k => k -> Histogram k -> Histogram k
increment Histogram k
forall a. Monoid a => a
mempty

fromCountList :: Ord k => [(k,Int)] -> Histogram k
fromCountList :: [(k, Int)] -> Histogram k
fromCountList = ((k, Int) -> Histogram k) -> [(k, Int)] -> Histogram k
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ((k -> Int -> Histogram k) -> (k, Int) -> Histogram k
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry k -> Int -> Histogram k
forall k. Ord k => k -> Int -> Histogram k
singletonCount)

flatMap :: Ord k' => (k -> Int -> Histogram k') -> Histogram k -> Histogram k'
flatMap :: (k -> Int -> Histogram k') -> Histogram k -> Histogram k'
flatMap k -> Int -> Histogram k'
f = ((k, Int) -> Histogram k') -> [(k, Int)] -> Histogram k'
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ((k -> Int -> Histogram k') -> (k, Int) -> Histogram k'
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry k -> Int -> Histogram k'
f) ([(k, Int)] -> Histogram k')
-> (Histogram k -> [(k, Int)]) -> Histogram k -> Histogram k'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Histogram k -> [(k, Int)]
forall k. Histogram k -> [(k, Int)]
toList

toList :: Histogram k -> [(k, Int)]
toList :: Histogram k -> [(k, Int)]
toList = Map k Int -> [(k, Int)]
forall k a. Map k a -> [(k, a)]
M.toList (Map k Int -> [(k, Int)])
-> (Histogram k -> Map k Int) -> Histogram k -> [(k, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Histogram k -> Map k Int
forall k. Histogram k -> Map k Int
toMap

lookup :: Ord k => k -> Histogram k -> Int
lookup :: k -> Histogram k -> Int
lookup k
k (Histogram Map k Int
m) = Int -> Maybe Int -> Int
forall a. a -> Maybe a -> a
fromMaybe Int
0 (Map k Int
m Map k Int -> k -> Maybe Int
forall k a. Ord k => Map k a -> k -> Maybe a
M.!? k
k)

-- | /O(n)/. The expression (@'split' k hist@) is a pair @(h1,h2)@
-- where all keys in @h1@ are lower than @k@ and all keys in
-- @h2@ larger than @k@. Any key equal to @k@ is found in neither @h1@ nor @h2@.
split :: Ord k => k -> Histogram k -> (Histogram k, Histogram k)
split :: k -> Histogram k -> (Histogram k, Histogram k)
split k
k (Histogram Map k Int
m) = let (Map k Int
lt, Map k Int
gt) = k -> Map k Int -> (Map k Int, Map k Int)
forall k a. Ord k => k -> Map k a -> (Map k a, Map k a)
M.split k
k Map k Int
m in (Map k Int -> Histogram k
forall k. Map k Int -> Histogram k
Histogram Map k Int
lt, Map k Int -> Histogram k
forall k. Map k Int -> Histogram k
Histogram Map k Int
gt)

splitLookup :: Ord k => k -> Histogram k -> (Histogram k, Int, Histogram k)
splitLookup :: k -> Histogram k -> (Histogram k, Int, Histogram k)
splitLookup k
k (Histogram Map k Int
m) = let (Map k Int
lt, Maybe Int
c, Map k Int
gt) = k -> Map k Int -> (Map k Int, Maybe Int, Map k Int)
forall k a. Ord k => k -> Map k a -> (Map k a, Maybe a, Map k a)
M.splitLookup k
k Map k Int
m in (Map k Int -> Histogram k
forall k. Map k Int -> Histogram k
Histogram Map k Int
lt, Int -> Maybe Int -> Int
forall a. a -> Maybe a -> a
fromMaybe Int
0 Maybe Int
c, Map k Int -> Histogram k
forall k. Map k Int -> Histogram k
Histogram Map k Int
gt)

(!) :: Ord k => Histogram k -> k -> Int
(!) = (k -> Histogram k -> Int) -> Histogram k -> k -> Int
forall a b c. (a -> b -> c) -> b -> a -> c
flip k -> Histogram k -> Int
forall k. Ord k => k -> Histogram k -> Int
Data.Histogram.lookup

-- | @'disjoint' k1 k2@ returns @True@ when there is no key that is nonzero in both @k1@ and @k2@.
disjoint :: Ord k => Histogram k -> Histogram k -> Bool

#if MIN_VERSION_containers (0,6,2)
disjoint :: Histogram k -> Histogram k -> Bool
disjoint (Histogram Map k Int
m1) (Histogram Map k Int
m2) = Map k Int -> Map k Int -> Bool
forall k a b. Ord k => Map k a -> Map k b -> Bool
M.disjoint Map k Int
m1 Map k Int
m2
#else
disjoint (Histogram m1) (Histogram m2) = M.null (M.intersection m1 m2)
#endif