--------------------------------------------------------------------------------
-- |
-- Module      :  Data.SignedMultiset
-- Copyright   :  (c) 2012 Stefan Holdermans
-- License     :  BSD-style
-- Maintainer  :  stefan@vectorfabrics.com
-- Stability   :  provisional
-- Portability :  portable
--
-- An efficient implementation of signed multisets.
--
-- A signed multiset is like a multiset (or bag), but additionally allows for
-- /negative membership/.
-- That is, in a signed multiset, an element can occur a negative number of
-- times.
--
-- For a theory of signed multisets, see
--
-- * Wayne D. Blizard. Negative membership.
--   /Notre Dame Journal of Formal Logic/, 31(3):346--368, 1990.
--
-- Since many function names (but not the type name) clash with Prelude names,
-- this module is usually imported @qualified@, e.g.,
--
-- >  import Data.SignedMultiset (SignedMultiset)
-- >  import qualified Data.SignedMultiset as SignedMultiset
--
-- Function comments contain the function's time complexity in so-called big-O
-- notation, with /n/ referring to the number of multiset members involved.
--
--------------------------------------------------------------------------------

module Data.SignedMultiset (

    -- * Type
    SignedMultiset,         -- abstract, instances: Eq, Ord, Show, Read, Monoid,
                            --   Typeable1, Data

    -- * Construction
    empty,                  -- :: SignedMultiset a
    singleton,              -- :: a -> SignedMultiset a
    insert,                 -- :: Ord a =>
                            --    a -> SignedMultiset a -> SignedMultiset a
    insertMany,             -- :: Ord a =>
                            --    a -> Int -> SignedMultiset a ->
                            --    SignedMultiset a
    delete,                 -- :: Ord a =>
                            --    a -> SignedMultiset a -> SignedMultiset a
    deleteMany,             -- :: Ord a =>
                            --    a -> Int -> SignedMultiset a ->
                            --    SignedMultiset a
    deleteAll,              -- :: Ord a =>
                            --    a -> SignedMultiset a -> SignedMultiset a

    -- * Queries
    null,                   -- :: SignedMultiset a -> Bool
    isSet,                  -- :: SignedMultiset a -> Bool
    size,                   -- :: SignedMultiset a -> Int
    cardinality,            -- :: SignedMultiset a -> Int
    member,                 -- :: Ord a => a -> SignedMultiset a -> Bool
    notMember,              -- :: Ord a => a -> SignedMultiset a -> Bool
    multiplicity,           -- :: Ord a => a -> SignedMultiset a -> Int
    isSubmultisetOf,        -- :: Ord a =>
                            --    SignedMultiset a -> SignedMultiset a -> Bool
    isProperSubmultisetOf,  -- :: Ord a =>
                            --    SignedMultiset a -> SignedMultiset a -> Bool

    -- * Combining
    union,                  -- :: Ord a =>
                            --    SignedMultiset a -> SignedMultiset a ->
                            --    SignedMultiset a
    additiveUnion,          -- :: Ord a =>
                            --    SignedMultiset a -> SignedMultiset a ->
                            --    SignedMultiset a
    intersection,           -- :: Ord a =>
                            --    SignedMultiset a -> SignedMultiset a ->
                            --    SignedMultiset a
    difference,             -- :: Ord a =>
                            --    SignedMultiset a -> SignedMultiset a ->
                            --    SignedMultiset a

    -- * Scalar multiplication
    multiply,               -- :: Ord a =>
                            --    Int -> SignedMultiset a -> SignedMultiset a

    -- * Traversals
    map,                    -- :: (Ord a, Ord b) =>
                            --    (a -> b) ->
                            --    SignedMultiset a -> SignedMultiset b
    foldr,                  -- :: (a -> b -> b) -> b -> SignedMultiset a -> b
    foldl,                  -- :: (a -> b -> a) -> a -> SignedMultiset b -> a

    -- * Conversion
    toList,                 -- :: SignedMultiset a -> [(a, Int)]
    fromList,               -- :: Ord a => [(a, Int)] -> SignedMultiset a

    -- * Additive wrapper
    Additive (..)           -- instances: Eq, Ord, Show, Read, Monoid

  ) where

import Prelude hiding (null, map, foldr, foldl)
import Data.Function (on)
import Data.Monoid (Monoid (..))
import Data.Data (Data (..), mkNoRepType)
import Data.Typeable (Typeable1 (..), mkTyCon3, mkTyConApp)
import Data.Map (Map)
import qualified Data.Map as Map
import Data.SignedMultiset.Show (showsMembers)
import Data.SignedMultiset.Read (readsMembers, mapReadS)

--------------------------------------------------------------------------------
-- Type
--------------------------------------------------------------------------------

-- Signed multisets are implemented as maps from elements to their
-- multiplicities.
-- We maintain the invariant that all elements are mapped to nonzero
-- multiplicities; i.e., if an element is not a member of the multiset, it does
-- not appear as a key in the corresponding map.

-- | A signed multiset with elements of type @a@.
newtype SignedMultiset a = SMS {unSMS :: Map a Int}

-- /O(n)/.
norm :: Ord a => Map a Int -> SignedMultiset a
norm = SMS . Map.filter (/= 0)

instance Ord a => Eq (SignedMultiset a) where
  (==) = (==) `on` unSMS

instance Ord a => Ord (SignedMultiset a) where
  compare = compare `on` unSMS

instance Show a => Show (SignedMultiset a) where
  showsPrec _ = showsMembers . toList

instance (Ord a, Read a) => Read (SignedMultiset a) where
  readsPrec = mapReadS fromList . readsMembers

-- | Monoid under 'union'.
instance Ord a => Monoid (SignedMultiset a) where
  mempty  = empty
  mappend = union

instance Typeable1 SignedMultiset where
  typeOf1 _ = mkTyConApp tyCon []
    where
      tyCon = mkTyCon3 "signed-multiset" "Data.SignedMultiset" "SignedMultiset"

instance (Ord a, Data a) => Data (SignedMultiset a) where
  gfoldl f z   = f (z fromList) . toList
  gunfold _    = error "Data.Data.gunfold: abstract datatype"
  toConstr _   = error "Data.Data.toConstr: abstract datatype"
  dataTypeOf _ = mkNoRepType "Data.SignedMultiset.SignedMultiset"

--------------------------------------------------------------------------------
-- Construction
--------------------------------------------------------------------------------

-- | /O(1)/. The empty signed multiset, i.e., the multiset in which every
-- element has multiplicity zero.
empty :: SignedMultiset a
empty = SMS Map.empty

-- | /O(1)/. Create a signed multiset that contains exactly one copy of the
-- given element.
singleton :: a -> SignedMultiset a
singleton x = SMS (Map.singleton x 1)

-- | /O(log n)/. Insert a new copy of the given element into a signed multiset,
-- i.e., increment the multiplicity of the element by 1.
insert :: Ord a => a -> SignedMultiset a -> SignedMultiset a
insert x = insertMany x 1

-- | /O(log n)/. Insert a specified number of new copies of the given element
-- into a signed multiset, i.e., increment the multiplicity of the element by
-- the specified number. If the specified number is negative, copies are
-- deleted from the set.
insertMany :: Ord a => a -> Int -> SignedMultiset a -> SignedMultiset a
insertMany x n = SMS . Map.alter f x . unSMS
  where
    f Nothing  = Just n
    f (Just m) = let k = m + n in if k == 0 then Nothing else Just k

-- | /O(log n)/. Delete a copy of the given element from a signed multiset,
-- i.e., decrement the multiplicity of the element by 1.
delete :: Ord a => a -> SignedMultiset a -> SignedMultiset a
delete x = deleteMany x 1

-- | /O(log n)/. Delete a specified number of copies of the given element from
-- a signed multiset, i.e., decrement the multiplicity of the element by the
-- specified number. If the specified number is negative, new copies of the
-- element are inserted into the set.
deleteMany :: Ord a => a -> Int -> SignedMultiset a -> SignedMultiset a
deleteMany x n = insertMany x (- n)

-- | /O(log n)/. Delete all copies of the given element from a signed multiset,
-- i.e., set the multiplicity of the element to zero.
deleteAll :: Ord a => a -> SignedMultiset a -> SignedMultiset a
deleteAll x = SMS . Map.delete x . unSMS

--------------------------------------------------------------------------------
-- Queries
--------------------------------------------------------------------------------

-- | /O(1)/. Return whether the signed multiset is empty, i.e., whether every
-- element has multiplicity zero.
null :: SignedMultiset a -> Bool
null = Map.null . unSMS

-- | /O(n)/. Return whether the signed multiset is a set, i.e., whether all
-- elements have either multiplicity zero or else multiplicity 1.
isSet :: SignedMultiset a -> Bool
isSet = Map.foldr ((&&) . (== 1)) True . unSMS

-- | /O(1)/. Return the number of members of the signed multiset, i.e., the
-- number of elements that have nonzero multiplicity.
size :: SignedMultiset a -> Int
size = Map.size . unSMS

-- | /O(n)/. Return the cardinality of the signed multiset, i.e., the sum of
-- the multiplicities of all elements.
cardinality :: SignedMultiset a -> Int
cardinality = Map.foldl' (+) 0 . unSMS

-- | /O(log n)/. Return whether the given element is a member of the signed
-- multiset, i.e., whether the element has nonzero multiplicity.
member :: Ord a => a -> SignedMultiset a -> Bool
member x = Map.member x . unSMS

-- | /O(log n)/. Return whether the given element is /not/ a member of the
-- signed multiset, i.e., whether the element has multiplicity zero.
notMember :: Ord a => a -> SignedMultiset a -> Bool
notMember x = Map.notMember x . unSMS

-- | /O(log n)/. Return the multiplicity of the given element in the signed
-- multiset.
multiplicity :: Ord a => a -> SignedMultiset a -> Int
multiplicity x = Map.findWithDefault 0 x . unSMS

-- | /O(n)/. Return whether the first signed multiset is a submultiset of the
-- second, i.e., whether each element that has nonzero multiplicity @n@ in the
-- first multiset has nonzero multiplicity @m@ with @n <= m@ in the second.
isSubmultisetOf :: Ord a => SignedMultiset a -> SignedMultiset a -> Bool
isSubmultisetOf = Map.isSubmapOfBy (<=) `on` unSMS

-- | /O(n)/. Return whether the first signed multiset is a proper
-- submultiset of the second, i.e., whether each element that has nonzero
-- multiplicity @n@ in the first multiset has nonzero multiplicity @m@ with
-- @n < m@ in the second.
isProperSubmultisetOf :: Ord a => SignedMultiset a -> SignedMultiset a -> Bool
isProperSubmultisetOf = Map.isSubmapOfBy (<) `on` unSMS

--------------------------------------------------------------------------------
-- Scalar multiplication
--------------------------------------------------------------------------------

-- | /O(n)/. Return the additive union of the given number of copies of the
-- signed multiset.
multiply :: Int -> SignedMultiset a -> SignedMultiset a
multiply n = SMS . Map.map (n *) . unSMS

--------------------------------------------------------------------------------
-- Combining
--------------------------------------------------------------------------------

-- | /O(n)/. Return the union of two signed multisets. The multiplicity of an
-- element in the returned multiset is the maximum of its nonzero
-- multiplicites in the argument multisets.
union :: Ord a => SignedMultiset a -> SignedMultiset a -> SignedMultiset a
union = SMS `after` (Map.unionWith max `on` unSMS)

-- | /O(n)/. Return the additive union of two signed multisets. The
-- multiplicity of an element in the returned multiset is the sum of its
-- multiplicities in the argument multisets.
additiveUnion :: Ord a =>
                 SignedMultiset a -> SignedMultiset a -> SignedMultiset a
additiveUnion = norm `after` (Map.unionWith (+) `on` unSMS)

-- | /O(n)/. Return the intersection of two signed multisets. If an element has
-- nonzero multiplicity in both argument multisets, its multiplicity in the
-- returned multiset is the minimum of its multiplicites in the argument
-- multisets; otherwise, its multiplicity in the returned multiset is zero.
intersection :: Ord a =>
                SignedMultiset a -> SignedMultiset a -> SignedMultiset a
intersection = SMS `after` (Map.intersectionWith min `on` unSMS)

-- | /O(n)/. Return the difference of two signed multisets. The multiplicity of
-- an element in the returned multiset is the difference between its
-- multiplicities in the first and second argument multiplicities.
difference :: Ord a => SignedMultiset a -> SignedMultiset a -> SignedMultiset a
difference = norm `after` (Map.unionWith (-) `on` unSMS)

--------------------------------------------------------------------------------
-- Traversals
--------------------------------------------------------------------------------

-- | /O(n * log n)/. Apply the given function to all elements of the signed
-- multiset.
map :: (Ord a, Ord b) => (a -> b) -> SignedMultiset a -> SignedMultiset b
map f = SMS . Map.mapKeys f . unSMS

-- | /O(n)/. Perform a right-associative fold on the members and elements of
-- the signed multiset using the given operator and start value.
foldr :: (a -> Int -> b -> b) -> b -> SignedMultiset a -> b
foldr f z = Map.foldrWithKey f z . unSMS

-- | /O(n)/. Perform a left-associative fold on the members and elements of
-- the signed multiset using the given operator and start value.
foldl :: (a -> b -> Int -> a) -> a -> SignedMultiset b -> a
foldl f z = Map.foldlWithKey f z . unSMS

--------------------------------------------------------------------------------
-- Conversion
--------------------------------------------------------------------------------

-- | /O(n)/. Convert the signed multiset to a list that associates all members
-- of the multiset with their multiplicity.
toList :: SignedMultiset a -> [(a, Int)]
toList = Map.toList . unSMS

-- | /O(n * log n)/. Construct a signed multiset from a list of
-- element/multiplicity pairs.
fromList :: Ord a => [(a, Int)] -> SignedMultiset a
fromList = SMS . Map.fromList . filter ((/= 0) . snd)

--------------------------------------------------------------------------------
-- Additive wrapper
--------------------------------------------------------------------------------

-- | Monoid under 'additiveUnion'.
newtype Additive a = Additive {getAdditive :: SignedMultiset a}
  deriving (Eq, Ord, Show, Read)

instance Ord a => Monoid (Additive a) where
  mempty  = Additive empty
  mappend = Additive `after` (additiveUnion `on` getAdditive)

--------------------------------------------------------------------------------
-- Utilities
--------------------------------------------------------------------------------

after :: (c -> d) -> (a -> b -> c) -> a -> b -> d
after f (.+.) x y = f (x .+. y)