--------------------------------------------------------------------------------
-- |
-- Module      :  Data.SignedMultiset
-- Copyright   :  (c) 2012 Stefan Holdermans
-- License     :  BSD-style
-- Maintainer  :  stefan@vectorfabrics.com
-- Stability   :  provisional
-- Portability :  non-portable (DeriveDataTypeable)
--
-- 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)}\"@.
--
--------------------------------------------------------------------------------

{-# LANGUAGE DeriveDataTypeable #-}

module Data.SignedMultiset (

    -- * Type
    SignedMultiset,         -- :: * -> *, abstract, instances: Eq, Ord, Show,
                            --   Read, Semigroup, 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,
                            --   Semigroup, 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.Data (Data (..), mkNoRepType)
import Data.Typeable (Typeable)
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} deriving Typeable

-- /O(n)/.
norm :: 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

-- | Semigroup under 'union'.
instance Ord a => Semigroup (SignedMultiset a) where
  (<>) = union

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

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 :: 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 :: 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 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 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 :: (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 :: (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 :: 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 z
  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 = norm . 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)

-- | Semigroup under 'additiveUnion'.
instance Ord a => Semigroup (Additive a) where
  (<>) = Additive `after` (additiveUnion `on` getAdditive)

-- | Monoid under 'additiveUnion'.
instance Ord a => Monoid (Additive a) where
  mempty  = Additive empty
  mappend = (<>)

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

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