{-# LANGUAGE DeriveAnyClass               #-}
{-# LANGUAGE NoGeneralisedNewtypeDeriving #-}
{-# LANGUAGE TypeApplications             #-}

module ZkFold.Base.Algebra.Polynomials.Multivariate.Monomial where

import           Control.DeepSeq                 (NFData)
import           Data.Aeson                      (FromJSON, ToJSON)
import           Data.List                       (intercalate)
import           Data.Map.Strict                 (Map, differenceWith, empty, filter, foldrWithKey, fromListWith,
                                                  intersectionWith, isSubmapOfBy, lookup, mapKeys, unionWith)
import qualified Data.Map.Strict                 as Map
import           GHC.Generics                    (Generic)
import           GHC.IsList                      (IsList (..))
import           Numeric.Natural                 (Natural)
import           Prelude                         hiding (Num (..), drop, filter, lcm, length, lookup, sum, take, (!!),
                                                  (/), (^))
import           Test.QuickCheck                 (Arbitrary (..))

import           ZkFold.Base.Algebra.Basic.Class

type Variable i = (Eq i, Ord i)

type Monomial i j = (Variable i, Ord j, Semiring j)

-- | Monomial type
newtype Mono i j = M (Map i j)
    deriving ((forall x. Mono i j -> Rep (Mono i j) x)
-> (forall x. Rep (Mono i j) x -> Mono i j) -> Generic (Mono i j)
forall x. Rep (Mono i j) x -> Mono i j
forall x. Mono i j -> Rep (Mono i j) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall i j x. Rep (Mono i j) x -> Mono i j
forall i j x. Mono i j -> Rep (Mono i j) x
$cfrom :: forall i j x. Mono i j -> Rep (Mono i j) x
from :: forall x. Mono i j -> Rep (Mono i j) x
$cto :: forall i j x. Rep (Mono i j) x -> Mono i j
to :: forall x. Rep (Mono i j) x -> Mono i j
Generic, Mono i j -> ()
(Mono i j -> ()) -> NFData (Mono i j)
forall a. (a -> ()) -> NFData a
forall i j. (NFData i, NFData j) => Mono i j -> ()
$crnf :: forall i j. (NFData i, NFData j) => Mono i j -> ()
rnf :: Mono i j -> ()
NFData, Value -> Parser [Mono i j]
Value -> Parser (Mono i j)
(Value -> Parser (Mono i j))
-> (Value -> Parser [Mono i j]) -> FromJSON (Mono i j)
forall a.
(Value -> Parser a) -> (Value -> Parser [a]) -> FromJSON a
forall i j.
(FromJSONKey i, Ord i, FromJSON j) =>
Value -> Parser [Mono i j]
forall i j.
(FromJSONKey i, Ord i, FromJSON j) =>
Value -> Parser (Mono i j)
$cparseJSON :: forall i j.
(FromJSONKey i, Ord i, FromJSON j) =>
Value -> Parser (Mono i j)
parseJSON :: Value -> Parser (Mono i j)
$cparseJSONList :: forall i j.
(FromJSONKey i, Ord i, FromJSON j) =>
Value -> Parser [Mono i j]
parseJSONList :: Value -> Parser [Mono i j]
FromJSON, [Mono i j] -> Value
[Mono i j] -> Encoding
Mono i j -> Value
Mono i j -> Encoding
(Mono i j -> Value)
-> (Mono i j -> Encoding)
-> ([Mono i j] -> Value)
-> ([Mono i j] -> Encoding)
-> ToJSON (Mono i j)
forall a.
(a -> Value)
-> (a -> Encoding)
-> ([a] -> Value)
-> ([a] -> Encoding)
-> ToJSON a
forall i j. (ToJSON j, ToJSONKey i) => [Mono i j] -> Value
forall i j. (ToJSON j, ToJSONKey i) => [Mono i j] -> Encoding
forall i j. (ToJSON j, ToJSONKey i) => Mono i j -> Value
forall i j. (ToJSON j, ToJSONKey i) => Mono i j -> Encoding
$ctoJSON :: forall i j. (ToJSON j, ToJSONKey i) => Mono i j -> Value
toJSON :: Mono i j -> Value
$ctoEncoding :: forall i j. (ToJSON j, ToJSONKey i) => Mono i j -> Encoding
toEncoding :: Mono i j -> Encoding
$ctoJSONList :: forall i j. (ToJSON j, ToJSONKey i) => [Mono i j] -> Value
toJSONList :: [Mono i j] -> Value
$ctoEncodingList :: forall i j. (ToJSON j, ToJSONKey i) => [Mono i j] -> Encoding
toEncodingList :: [Mono i j] -> Encoding
ToJSON)

------------------------------------ Map-based monomials ------------------------------------

-- | Monomial constructor
monomial :: Monomial i j => Map i j -> Mono i j
monomial :: forall i j. Monomial i j => Map i j -> Mono i j
monomial = Map i j -> Mono i j
forall i j. Map i j -> Mono i j
M (Map i j -> Mono i j)
-> (Map i j -> Map i j) -> Map i j -> Mono i j
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (j -> Bool) -> Map i j -> Map i j
forall a k. (a -> Bool) -> Map k a -> Map k a
filter (j -> j -> Bool
forall a. Eq a => a -> a -> Bool
/= j
forall a. AdditiveMonoid a => a
zero)

evalMonomial :: forall i j b .
    MultiplicativeMonoid b =>
    Exponent b j =>
    (i -> b) -> Mono i j -> b
evalMonomial :: forall i j b.
(MultiplicativeMonoid b, Exponent b j) =>
(i -> b) -> Mono i j -> b
evalMonomial i -> b
f (M Map i j
m) =
    (i -> j -> b -> b) -> b -> Map i j -> b
forall k a b. (k -> a -> b -> b) -> b -> Map k a -> b
foldrWithKey (\i
i j
j b
x -> (i -> b
f i
i b -> j -> b
forall a b. Exponent a b => a -> b -> a
^ j
j) b -> b -> b
forall a. MultiplicativeSemigroup a => a -> a -> a
* b
x) (forall a. MultiplicativeMonoid a => a
one @b) Map i j
m

-- | Maps a variable index using the provided `Map`
mapVar :: Variable i => Map i i -> i -> i
mapVar :: forall i. Variable i => Map i i -> i -> i
mapVar Map i i
m i
x = case i
x i -> Map i i -> Maybe i
forall k a. Ord k => k -> Map k a -> Maybe a
`lookup` Map i i
m of
    Just i
y -> i
y
    Maybe i
_      -> [Char] -> i
forall a. HasCallStack => [Char] -> a
error [Char]
"mapVar: something went wrong"

mapVarMonomial :: Variable i => Map i i -> Mono i j -> Mono i j
mapVarMonomial :: forall i j. Variable i => Map i i -> Mono i j -> Mono i j
mapVarMonomial Map i i
m (M Map i j
as) = Map i j -> Mono i j
forall i j. Map i j -> Mono i j
M (Map i j -> Mono i j) -> Map i j -> Mono i j
forall a b. (a -> b) -> a -> b
$ (i -> i) -> Map i j -> Map i j
forall k2 k1 a. Ord k2 => (k1 -> k2) -> Map k1 a -> Map k2 a
mapKeys (Map i i -> i -> i
forall i. Variable i => Map i i -> i -> i
mapVar Map i i
m) Map i j
as

instance Monomial i j => IsList (Mono i j) where
    type Item (Mono i j) = (i, j)
    toList :: Mono i j -> [Item (Mono i j)]
toList (M Map i j
m) = Map i j -> [Item (Map i j)]
forall l. IsList l => l -> [Item l]
toList Map i j
m
    fromList :: [Item (Mono i j)] -> Mono i j
fromList [Item (Mono i j)]
m = Map i j -> Mono i j
forall i j. Map i j -> Mono i j
M (Map i j -> Mono i j) -> Map i j -> Mono i j
forall a b. (a -> b) -> a -> b
$ (j -> j -> j) -> [(i, j)] -> Map i j
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
fromListWith j -> j -> j
forall a. AdditiveSemigroup a => a -> a -> a
(+) [(i, j)]
[Item (Mono i j)]
m

instance (Show i, Show j, Monomial i j) => Show (Mono i j) where
    show :: Mono i j -> [Char]
show (M Map i j
m) = [Char] -> [[Char]] -> [Char]
forall a. [a] -> [[a]] -> [a]
intercalate [Char]
"∙" ([[Char]] -> [Char])
-> ([(i, j)] -> [[Char]]) -> [(i, j)] -> [Char]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((i, j) -> [Char]) -> [(i, j)] -> [[Char]]
forall a b. (a -> b) -> [a] -> [b]
map (i, j) -> [Char]
showVar ([(i, j)] -> [Char]) -> [(i, j)] -> [Char]
forall a b. (a -> b) -> a -> b
$ Map i j -> [Item (Map i j)]
forall l. IsList l => l -> [Item l]
toList Map i j
m
        where
            showVar :: (i, j) -> String
            showVar :: (i, j) -> [Char]
showVar (i
i, j
j) = [Char]
"x" [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ i -> [Char]
forall a. Show a => a -> [Char]
show i
i [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ (if j
j j -> j -> Bool
forall a. Eq a => a -> a -> Bool
== j
forall a. MultiplicativeMonoid a => a
one then [Char]
"" else [Char]
"^" [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ j -> [Char]
forall a. Show a => a -> [Char]
show j
j)

instance Monomial i j => Eq (Mono i j) where
    M Map i j
asl == :: Mono i j -> Mono i j -> Bool
== M Map i j
asr = Map i j
asl Map i j -> Map i j -> Bool
forall a. Eq a => a -> a -> Bool
== Map i j
asr

instance Monomial i j => Ord (Mono i j) where
    compare :: Mono i j -> Mono i j -> Ordering
compare (M Map i j
asl) (M Map i j
asr) = [(i, j)] -> [(i, j)] -> Ordering
forall {a} {a}. (Ord a, Ord a) => [(a, a)] -> [(a, a)] -> Ordering
go (Map i j -> [Item (Map i j)]
forall l. IsList l => l -> [Item l]
toList Map i j
asl) (Map i j -> [Item (Map i j)]
forall l. IsList l => l -> [Item l]
toList Map i j
asr)
        where
            go :: [(a, a)] -> [(a, a)] -> Ordering
go [] [] = Ordering
EQ
            go [] [(a, a)]
_  = Ordering
LT
            go [(a, a)]
_  [] = Ordering
GT
            go ((a
k1, a
a1):[(a, a)]
xs) ((a
k2, a
a2):[(a, a)]
ys)
                | a
k1 a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
k2  = if a
a1 a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
a2 then [(a, a)] -> [(a, a)] -> Ordering
go [(a, a)]
xs [(a, a)]
ys else a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare a
a1 a
a2
                | Bool
otherwise = a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare a
k2 a
k1

instance (Monomial i j, Arbitrary i, Arbitrary j) => Arbitrary (Mono i j) where
    arbitrary :: Gen (Mono i j)
arbitrary = Map i j -> Mono i j
forall i j. Map i j -> Mono i j
M (Map i j -> Mono i j) -> Gen (Map i j) -> Gen (Mono i j)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen (Map i j)
forall a. Arbitrary a => Gen a
arbitrary

instance Monomial i j => MultiplicativeSemigroup (Mono i j) where
    M Map i j
l * :: Mono i j -> Mono i j -> Mono i j
* M Map i j
r = Map i j -> Mono i j
forall i j. Map i j -> Mono i j
M (Map i j -> Mono i j) -> Map i j -> Mono i j
forall a b. (a -> b) -> a -> b
$ (j -> Bool) -> Map i j -> Map i j
forall a k. (a -> Bool) -> Map k a -> Map k a
Map.filter (j -> j -> Bool
forall a. Eq a => a -> a -> Bool
/= j
forall a. AdditiveMonoid a => a
zero) (Map i j -> Map i j) -> Map i j -> Map i j
forall a b. (a -> b) -> a -> b
$ (j -> j -> j) -> Map i j -> Map i j -> Map i j
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
unionWith j -> j -> j
forall a. AdditiveSemigroup a => a -> a -> a
(+) Map i j
l Map i j
r

instance Monomial i j => Exponent (Mono i j) Natural where
    ^ :: Mono i j -> Natural -> Mono i j
(^) = Mono i j -> Natural -> Mono i j
forall a. MultiplicativeMonoid a => a -> Natural -> a
natPow

instance Monomial i j => MultiplicativeMonoid (Mono i j) where
    one :: Mono i j
one = Map i j -> Mono i j
forall i j. Map i j -> Mono i j
M Map i j
forall k a. Map k a
empty

instance (Monomial i j, Ring j) => Exponent (Mono i j) Integer where
    ^ :: Mono i j -> Integer -> Mono i j
(^) = Mono i j -> Integer -> Mono i j
forall a. MultiplicativeGroup a => a -> Integer -> a
intPow

instance (Monomial i j, Ring j) => MultiplicativeGroup (Mono i j) where
    invert :: Mono i j -> Mono i j
invert (M Map i j
m) = Map i j -> Mono i j
forall i j. Map i j -> Mono i j
M (Map i j -> Mono i j) -> Map i j -> Mono i j
forall a b. (a -> b) -> a -> b
$ (j -> j) -> Map i j -> Map i j
forall a b k. (a -> b) -> Map k a -> Map k b
Map.map j -> j
forall a. AdditiveGroup a => a -> a
negate (Map i j -> Map i j) -> Map i j -> Map i j
forall a b. (a -> b) -> a -> b
$ Map i j
m

    M Map i j
l / :: Mono i j -> Mono i j -> Mono i j
/ M Map i j
r = Map i j -> Mono i j
forall i j. Map i j -> Mono i j
M (Map i j -> Mono i j) -> Map i j -> Mono i j
forall a b. (a -> b) -> a -> b
$ (j -> j -> Maybe j) -> Map i j -> Map i j -> Map i j
forall k a b.
Ord k =>
(a -> b -> Maybe a) -> Map k a -> Map k b -> Map k a
differenceWith j -> j -> Maybe j
forall {a}. (Eq a, AdditiveGroup a) => a -> a -> Maybe a
f Map i j
l Map i j
r
        where f :: a -> a -> Maybe a
f a
a a
b = if a
a a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
b then Maybe a
forall a. Maybe a
Nothing else a -> Maybe a
forall a. a -> Maybe a
Just (a
a a -> a -> a
forall a. AdditiveGroup a => a -> a -> a
- a
b)

oneM :: Mono i j -> Bool
oneM :: forall i j. Mono i j -> Bool
oneM (M Map i j
m) = Map i j -> Bool
forall k a. Map k a -> Bool
Map.null Map i j
m

dividable :: forall i j . Monomial i j => Mono i j -> Mono i j -> Bool
dividable :: forall i j. Monomial i j => Mono i j -> Mono i j -> Bool
dividable (M Map i j
l) (M Map i j
r) = (j -> j -> Bool) -> Map i j -> Map i j -> Bool
forall k a b.
Ord k =>
(a -> b -> Bool) -> Map k a -> Map k b -> Bool
isSubmapOfBy j -> j -> Bool
forall a. Ord a => a -> a -> Bool
(<=) Map i j
r Map i j
l

lcmM :: Monomial i j => Mono i j -> Mono i j -> Mono i j
lcmM :: forall i j. Monomial i j => Mono i j -> Mono i j -> Mono i j
lcmM (M Map i j
l) (M Map i j
r) = Map i j -> Mono i j
forall i j. Map i j -> Mono i j
M (Map i j -> Mono i j) -> Map i j -> Mono i j
forall a b. (a -> b) -> a -> b
$ (j -> j -> j) -> Map i j -> Map i j -> Map i j
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
unionWith j -> j -> j
forall a. Ord a => a -> a -> a
max Map i j
l Map i j
r

gcdM :: Monomial i j => Mono i j -> Mono i j -> Mono i j
gcdM :: forall i j. Monomial i j => Mono i j -> Mono i j -> Mono i j
gcdM (M Map i j
l) (M Map i j
r) = Map i j -> Mono i j
forall i j. Map i j -> Mono i j
M ((j -> j -> j) -> Map i j -> Map i j -> Map i j
forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
intersectionWith j -> j -> j
forall a. Ord a => a -> a -> a
min Map i j
l Map i j
r)