--------------------------------------------------------------------------------
-- |
-- 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 object 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.
--
-- Signed-multiset types are constructed by the type constructor
-- 'SignedMultiset'.
-- The number of times an object appears in a signed multiset is called its
-- 'multiplicity'.
-- An object is said to be a 'member' of a signed multiset if it has a nonzero
-- multiplicity.
-- The number of members of a signed multiset is referred to as its 'size',
-- while the 'cardinality' of a signed multiset is the sum of the multiplicities
-- of its members.
-- A signed multiset is 'empty' if it is without members.
--
-- Textually, signed multisets are represented by listing their members and, in
-- parentheses, their multiplicities between curly brackets.
-- For instance, the signed multiset that contains -1 copies of 2, 2 copies of
-- 3, and -4 copies of 5 is denoted by @\"{2(-1),3(2),5(-4)}\"@.
--
--------------------------------------------------------------------------------

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
    isPositive,             -- :: SignedMultiset a -> Bool
    isNegative,             -- :: 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

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

    -- * Memberwise operations
    root,                   -- :: SignedMultiset a -> SignedMultiset a
    shadow,                 -- :: SignedMultiset a -> SignedMultiset a
    modulus,                -- :: SignedMultiset a -> SignedMultiset a
    signum,                 -- :: SignedMultiset a -> SignedMultiset a
    unitstep,               -- :: Ord a => SignedMultiset a -> SignedMultiset a
    multiply,               -- :: Ord a =>
                            --    Int -> SignedMultiset a -> SignedMultiset a

    -- * Traversals
    map,                    -- :: (Ord a, Ord b) =>
                            --    (a -> b) ->
                            --    SignedMultiset a -> SignedMultiset b
    additiveMap,            -- :: (Ord a, Ord b) =>
                            --    (a -> b) ->
                            --    SignedMultiset a -> SignedMultiset b
    filter,                 -- :: Ord a =>
                            --    (a -> Int -> Bool) -> SignedMultiset a ->
                            --    SignedMultiset a
    partition,              -- :: Ord a =>
                            --    (a -> Int -> Bool) -> SignedMultiset a ->
                            --    (SignedMultiset a, SignedMultiset a)
    split,                  -- :: Ord a =>
                            --    (a -> Int -> Bool) -> SignedMultiset a ->
                            --    (SignedMultiset a, SignedMultiset a)
    foldr,                  -- :: (a -> b -> b) -> b -> SignedMultiset a -> b
    foldr',                 -- :: (a -> b -> b) -> b -> SignedMultiset a -> b
    foldl,                  -- :: (a -> b -> a) -> a -> SignedMultiset b -> a
    foldl',                 -- :: (a -> b -> a) -> a -> SignedMultiset b -> a

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

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

  ) where

import Prelude hiding (signum, null, map, filter, foldr, foldl)
import qualified Prelude (signum, filter, foldr)
import Control.Arrow ((>>>), (***))
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 objects to their
-- multiplicities.
-- We maintain the invariant that all objects are mapped to nonzero
-- multiplicities; i.e., if an object is not a member of the multiset, it does
-- not appear as a key in the corresponding map.

-- | A signed multiset over objects 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 object
-- has multiplicity zero.
empty :: SignedMultiset a
empty = SMS Map.empty

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

-- | /O(log n)/. Insert a new copy of the given object into a signed multiset,
-- i.e., increment the multiplicity of the object 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 object
-- into a signed multiset, i.e., increment the multiplicity of the object 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 object from a signed multiset, i.e.,
-- decrement the multiplicity of the object 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 object from
-- a signed multiset, i.e., decrement the multiplicity of the object by the
-- specified number. If the specified number is negative, new copies of the
-- object 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 object from a signed multiset,
-- i.e., set the multiplicity of the object 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
-- object 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
-- object have either multiplicity zero or else multiplicity 1.
isSet :: SignedMultiset a -> Bool
isSet = Map.foldr ((&&) . (== 1)) True . unSMS

-- | /O(n)/. Return whether all objects in the signed multiset have nonnegative
-- multiplicities.
isPositive :: SignedMultiset a -> Bool
isPositive = Map.foldr ((&&) . (> 0)) True . unSMS

-- | /O(n)/. Return whether all objects in the signed multiset have nonpositive
-- multiplicities.
isNegative :: SignedMultiset a -> Bool
isNegative = Map.foldr ((&&) . (< 0)) True . unSMS

-- | /O(1)/. Return the number of members of the signed multiset, i.e., the
-- number of objects 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 objects.
cardinality :: SignedMultiset a -> Int
cardinality = Map.foldl' (+) 0 . unSMS

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

-- | /O(log n)/. Return whether the given object is /not/ a member of the signed
-- multiset, i.e., whether the object 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 object 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 object that has nonzero multiplicity @m@ in the
-- first multiset has nonzero multiplicity @n@ with @m <= n@ in the second.
isSubmultisetOf :: Ord a => SignedMultiset a -> SignedMultiset a -> Bool
isSubmultisetOf = Map.isSubmapOfBy (<=) `on` unSMS

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

-- | /O(n)/. Return the union of two signed multisets. The multiplicity of an
-- object in the returned multiset is the maximum of its nonzero multiplicities
-- 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 object 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 object has
-- nonzero multiplicity in both argument multisets, its multiplicity in the
-- returned multiset is the minimum of its multiplicities 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)

--------------------------------------------------------------------------------
-- Memberwise operations
--------------------------------------------------------------------------------

-- | /O(n)/. Return the root of the signed multiset. The multiplicity of an
-- object in the returned multiset is zero if its multiplicity in the argument
-- multiset is zero and 1 otherwise.
root :: SignedMultiset a -> SignedMultiset a
root = SMS . Map.map (const 1) . unSMS

-- | /O(n)/. Return the shadow of the signed multiset. The multiplicity of an
-- object in the returned multiset is the additive inverse of its multiplicity
-- in the argument multiset.
shadow :: SignedMultiset a -> SignedMultiset a
shadow = SMS . Map.map negate . unSMS

-- | /O(n)/. Return the modulus of the signed multiset. The multiplicity of an
-- object in the returned multiset is the absolute value of its multiplicity in
-- the argument multiset.
modulus :: SignedMultiset a -> SignedMultiset a
modulus = SMS . Map.map abs . unSMS

-- | /O(n)/. Return the signum of the signed multiset. The multiplicity of an
-- object in the returned multiset is -1 if it has negative multiplicity in the
-- argument multiset, zero if its multiplicity in the argument multiset is zero,
-- and 1 if it has positive multiplicity in the argument multiset.
signum :: SignedMultiset a -> SignedMultiset a
signum = SMS . Map.map Prelude.signum . unSMS

-- | /O(n)/. Return the left-continuous unit step of the signed multiset. The
-- multiplicity of an object in the returned multiset is zero if it has negative
-- multiplicity in the argument multiset, and 1 otherwise.
unitstep :: Ord a => SignedMultiset a -> SignedMultiset a
unitstep = norm . Map.map u . unSMS
  where
    u n = if n > 0 then 1 else 0

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

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

-- | /O(n * log n)/. Apply the given function to all objects of the signed
-- multiset. If the the function maps distinct objects to the same new object,
-- the multiplicity of the new object is the maximum of the nonzero 
-- multiplicities of the two original objects.
map :: (Ord a, Ord b) =>  (a -> b) -> SignedMultiset a -> SignedMultiset b
map f = SMS . Map.mapKeysWith max f . unSMS

-- | /O(n * log n)/. Apply the given function to all objects of the signed
-- multiset. If the the function maps distinct objects to the same new object,
-- the multiplicity of the new object is the sum of the multiplicities of the
-- two original objects.
additiveMap :: (Ord a, Ord b) =>
               (a -> b) -> SignedMultiset a -> SignedMultiset b
additiveMap f = norm . Map.mapKeysWith (+) f . unSMS

-- | /O(n)/. Apply the given predicate to the members of the signed multiset and
-- their multiplicities. The returned multiset contains the copies of the
-- members that satisfy the predicate.
filter :: Ord a => (a -> Int -> Bool) -> SignedMultiset a -> SignedMultiset a
filter p = SMS . Map.filterWithKey p . unSMS

-- | /O(n)/. Apply the given predicate to the members of the signed multiset and
-- their multiplicity. The first returned multiset contains the copies of the
-- members that satisfy the predicate, while the second returned multiset
-- contains the copies of the members that do not satisfy the predicate.
partition :: Ord a =>
             (a -> Int -> Bool) -> SignedMultiset a ->
             (SignedMultiset a, SignedMultiset a)
partition p = (SMS *** SMS) . Map.partitionWithKey p . unSMS

-- | /O(n)/. Split the signed multiset into a multiset containing the copies of
-- the members with a multiplicity less than or equal to the given number and a
-- multiset containing the copies of the members with a multiplicity greater
-- than the given number.
split :: Ord a =>
         Int -> SignedMultiset a -> (SignedMultiset a, SignedMultiset a)
split n = partition p
  where
    p = const (< n)

-- | /O(n)/. Perform a right-associative fold on the members of the signed
-- multiset and their multiplicities 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 strict right-associative fold on the members of the
-- signed multiset and their multiplicities 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 of the signed
-- multiset and their multiplicities using the given operator and start value.
foldl :: (a -> b -> Int -> a) -> a -> SignedMultiset b -> a
foldl f z = Map.foldlWithKey f z . unSMS

-- | /O(n)/. Perform a strict left-associative fold on the members of the signed
-- multiset and their multiplicities 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 + k)/ (with /k/ the combined length of the returned lists). Return two
-- lists, such that: for each object with a positive multiplicity @m@ in the
-- signed multiset, the first list contains @m@ copies and the second list
-- contains no copies of the object; for each object with a negative
-- multiplicity @-n@, the first list contains no and the second list contains
-- @n@ copies of the object; and for each object with zero multiplicity,
-- neither list contains a copy of the object.
toLists :: SignedMultiset a -> ([a], [a])
toLists = foldr f ([], [])
  where
    z              = ([], [])
    f x n (xs, ys) =
      if n > 0 then (replicate n x ++ xs, ys) else (xs, replicate (- n) x ++ ys)

-- | /O(k * log n)/ (with /k/ the length of the argument list). Construct a
-- signed multiset from a list of object/multiplicity pairs.
fromList :: Ord a => [(a, Int)] -> SignedMultiset a
fromList = SMS . Map.fromListWith (+) . Prelude.filter ((/= 0) . snd)

-- | /O(k * log n)/ (with /k/ the combined length of the argument lists).
-- Construct a signed multiset by, starting from the empty multiset, inserting
-- copies of objects from the first argument list and deleting copies of objects
-- from the second argument list.
fromLists :: Ord a => [a] -> [a] -> SignedMultiset a
fromLists = Prelude.foldr delete . Prelude.foldr insert empty

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

-- | An element of the free abelian group on @a@.
newtype Additive a = Additive {getAdditive :: SignedMultiset a}
  deriving (Eq, Ord, Show, Read)

-- | Monoid under 'additiveUnion'.
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)