{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE StandaloneDeriving #-}
module Autolib.Multiset where

import Autolib.ToDoc
import Autolib.Reader
import qualified Data.Map.Strict as M
import Control.Monad (mfilter)

newtype Multiset a = Multiset (M.Map a Int)

fromList :: Ord a => [a] -> Multiset a
fromList :: forall a. Ord a => [a] -> Multiset a
fromList [a]
vs = [(a, Int)] -> Multiset a
forall a. Ord a => [(a, Int)] -> Multiset a
multiset ([(a, Int)] -> Multiset a) -> [(a, Int)] -> Multiset a
forall a b. (a -> b) -> a -> b
$ [a] -> [Int] -> [(a, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [a]
vs ([Int] -> [(a, Int)]) -> [Int] -> [(a, Int)]
forall a b. (a -> b) -> a -> b
$ Int -> [Int]
forall a. a -> [a]
repeat Int
1

multiset :: Ord a => [(a,Int)] -> Multiset a
multiset :: forall a. Ord a => [(a, Int)] -> Multiset a
multiset [(a, Int)]
kvs = 
    Map a Int -> Multiset a
forall a. Map a Int -> Multiset a
Multiset (Map a Int -> Multiset a) -> Map a Int -> Multiset a
forall a b. (a -> b) -> a -> b
$ (Int -> Bool) -> Map a Int -> Map a Int
forall a k. (a -> Bool) -> Map k a -> Map k a
M.filter (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0) (Map a Int -> Map a Int) -> Map a Int -> Map a Int
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Int) -> [(a, Int)] -> Map a Int
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
M.fromListWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) [(a, Int)]
kvs

instance ToDoc a => ToDoc (Multiset a) where
    toDocPrec :: Int -> Multiset a -> Doc
toDocPrec Int
_ (Multiset Map a Int
m) = Doc -> Doc
braces 
        (Doc -> Doc) -> Doc -> Doc
forall a b. (a -> b) -> a -> b
$ [Doc] -> Doc
hsep ([Doc] -> Doc) -> [Doc] -> Doc
forall a b. (a -> b) -> a -> b
$ Doc -> [Doc] -> [Doc]
forall {t}. Semigroup t => t -> [t] -> [t]
punctuate Doc
comma
        ([Doc] -> [Doc]) -> [Doc] -> [Doc]
forall a b. (a -> b) -> a -> b
$ ((a, Int) -> Doc) -> [(a, Int)] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map (\(a
k,Int
v) -> [Doc] -> Doc
hcat [ a -> Doc
forall a. ToDoc a => a -> Doc
toDoc a
k, String -> Doc
text String
":", Int -> Doc
forall a. ToDoc a => a -> Doc
toDoc Int
v ])
        ([(a, Int)] -> [Doc]) -> [(a, Int)] -> [Doc]
forall a b. (a -> b) -> a -> b
$ Map a Int -> [(a, Int)]
forall k a. Map k a -> [(k, a)]
M.toList Map a Int
m

instance (Ord a, Reader a) => Reader (Multiset a) where
    atomic_readerPrec :: Int -> Parser (Multiset a)
atomic_readerPrec Int
_ = [(a, Int)] -> Multiset a
forall a. Ord a => [(a, Int)] -> Multiset a
multiset ([(a, Int)] -> Multiset a)
-> ParsecT String () Identity [(a, Int)] -> Parser (Multiset a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ( ParsecT String () Identity [(a, Int)]
-> ParsecT String () Identity [(a, Int)]
forall a. Parser a -> Parser a
my_braces (ParsecT String () Identity [(a, Int)]
 -> ParsecT String () Identity [(a, Int)])
-> ParsecT String () Identity [(a, Int)]
-> ParsecT String () Identity [(a, Int)]
forall a b. (a -> b) -> a -> b
$ Parser (a, Int) -> ParsecT String () Identity [(a, Int)]
forall a. Parser a -> Parser [a]
my_commaSep (Parser (a, Int) -> ParsecT String () Identity [(a, Int)])
-> Parser (a, Int) -> ParsecT String () Identity [(a, Int)]
forall a b. (a -> b) -> a -> b
$ do
        k <- Parser a
forall a. Reader a => Parser a
reader
        my_symbol ":"
        v <- mfilter (> 0) reader <|> fail "must be positive"
        return (k,v)
     )

deriving via FromReader (Multiset a) instance (Ord a, Reader a) => Read (Multiset a)

empty :: Multiset a
empty :: forall a. Multiset a
empty = (Map a Int -> Multiset a
forall a. Map a Int -> Multiset a
Multiset Map a Int
forall k a. Map k a
M.empty)

null :: Ord a => Multiset a -> Bool
null :: forall a. Ord a => Multiset a -> Bool
null (Multiset Map a Int
a) = Map a Int -> Bool
forall k a. Map k a -> Bool
M.null Map a Int
a

union :: Ord a => Multiset a -> Multiset a -> Multiset a
union :: forall a. Ord a => Multiset a -> Multiset a -> Multiset a
union (Multiset Map a Int
a) (Multiset Map a Int
b) = 
    Map a Int -> Multiset a
forall a. Map a Int -> Multiset a
Multiset (Map a Int -> Multiset a) -> Map a Int -> Multiset a
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Int) -> Map a Int -> Map a Int -> Map a Int
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
M.unionWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) Map a Int
a Map a Int
b

difference :: Ord a => Multiset a -> Multiset a -> Multiset a
difference :: forall a. Ord a => Multiset a -> Multiset a -> Multiset a
difference (Multiset Map a Int
a) (Multiset Map a Int
b) = 
    Map a Int -> Multiset a
forall a. Map a Int -> Multiset a
Multiset (Map a Int -> Multiset a) -> Map a Int -> Multiset a
forall a b. (a -> b) -> a -> b
$ (Int -> Bool) -> Map a Int -> Map a Int
forall a k. (a -> Bool) -> Map k a -> Map k a
M.filter (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) 
    (Map a Int -> Map a Int) -> Map a Int -> Map a Int
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Maybe Int) -> Map a Int -> Map a Int -> Map a Int
forall k a b.
Ord k =>
(a -> b -> Maybe a) -> Map k a -> Map k b -> Map k a
M.differenceWith ( \ Int
x Int
y -> Int -> Maybe Int
forall a. a -> Maybe a
Just (Int
xInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
y) ) Map a Int
a Map a Int
b

symmetric_difference :: Ord a => Multiset a -> Multiset a -> Multiset a
symmetric_difference :: forall a. Ord a => Multiset a -> Multiset a -> Multiset a
symmetric_difference (Multiset Map a Int
a) (Multiset Map a Int
b) = 
      Map a Int -> Multiset a
forall a. Map a Int -> Multiset a
Multiset (Map a Int -> Multiset a) -> Map a Int -> Multiset a
forall a b. (a -> b) -> a -> b
$ (Int -> Bool) -> Map a Int -> Map a Int
forall a k. (a -> Bool) -> Map k a -> Map k a
M.filter (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) 
    (Map a Int -> Map a Int) -> Map a Int -> Map a Int
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Int) -> Map a Int -> Map a Int -> Map a Int
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
M.unionWith ( \ Int
x Int
y -> Int -> Int
forall a. Num a => a -> a
abs (Int
xInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
y)) Map a Int
a Map a Int
b

intersection :: Ord a => Multiset a -> Multiset a -> Multiset a
intersection :: forall a. Ord a => Multiset a -> Multiset a -> Multiset a
intersection (Multiset Map a Int
a) (Multiset Map a Int
b) = 
    Map a Int -> Multiset a
forall a. Map a Int -> Multiset a
Multiset (Map a Int -> Multiset a) -> Map a Int -> Multiset a
forall a b. (a -> b) -> a -> b
$ (Int -> Bool) -> Map a Int -> Map a Int
forall a k. (a -> Bool) -> Map k a -> Map k a
M.filter (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (Map a Int -> Map a Int) -> Map a Int -> Map a Int
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Int) -> Map a Int -> Map a Int -> Map a Int
forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Map a Int
a Map a Int
b