module NLP.Scores.Internals
    ( Counts(..)
    , Count
    , P(..)
    , empty
    , unionPlus
    )
where
import qualified Data.Map as Map
import Data.Monoid

-- | A count
type Count = Double
-- | Count table
data Counts a b = 
  Counts 
  { joint :: !(Map.Map (P a b) Count) -- ^ Counts of both components
  , marginalFst :: !(Map.Map a Count) -- ^ Counts of the first component
  , marginalSnd :: !(Map.Map b Count) -- ^ Counts of the second component
  } 
data P a b = P !a !b deriving (Eq, Ord)

-- | The empty count table
empty :: (Ord a, Ord b) => Counts a b
empty = Counts Map.empty Map.empty Map.empty

instance (Ord a, Ord b) => Monoid (Counts a b) where
    mempty = empty
    c `mappend` k = 
        Counts { joint = unionPlus (joint c) (joint k)
               , marginalFst = unionPlus (marginalFst c) (marginalFst k)
               , marginalSnd = unionPlus (marginalSnd c) (marginalSnd k)
               }

unionPlus :: (Num a, Ord k) => Map.Map k a -> Map.Map k a -> Map.Map k a
unionPlus m = 
    Map.foldlWithKey' (\z k v -> Map.insertWith' (+) k v z) m
{-# SPECIALIZE unionPlus :: (Ord k) => 
  Map.Map k Count -> Map.Map k Count -> Map.Map k Count #-}