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

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

    -- * Traversals
    map,                    -- :: (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 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(n)/. Return whether all elements in the signed multiset have nonnegative
-- multiplicities.
isPositive :: SignedMultiset a -> Bool
isPositive = Map.foldr ((&&) . (> 0)) True . unSMS

-- | /O(n)/. Return whether all elements 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 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 @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

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

--------------------------------------------------------------------------------
-- Elementwise operations
--------------------------------------------------------------------------------

-- | /O(n)/. Return the shadow of the signed multiset. The multiplicity of an
-- element 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
-- element 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
-- element 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 element in the returned multiset is zero if it has
-- negative multiplicity in the argument multiset, and 1 otherwise.
unitstep :: SignedMultiset a -> SignedMultiset a
unitstep = SMS . 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 :: Int -> SignedMultiset a -> SignedMultiset a
multiply n = SMS . Map.map (n *) . 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)/. 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 element with a positive multiplicity @m@ in the
-- signed multiset, the first list contains @m@ copies and the second list
-- contains no copies of the element; for each element with a negative
-- multiplicity @- n@, the first list contains no and the second list contains
-- @n@ copies of the element; and for each element with zero multiplicity,
-- neither list contains a copy of the element.
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 element/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 elements from the first argument list and deleting copies of
-- elements 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)