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

module ZkFold.Base.Algebra.Polynomials.Multivariate.Monomial
    ( M(..)
    , Monomial
    , Monomial'
    , monomial
    , Variable
    , MonomialAny
    , MonomialRepAny
    , MonomialRepBoundedDegree
    , MonomialBoundedDegree
    , evalMapM
    , evalVectorM
    , mapVar
    , mapVarMonomial
    ) where

import           Control.DeepSeq                 (NFData)
import           Data.Aeson                      (FromJSON, ToJSON)
import           Data.List                       (intercalate)
import           Data.Map.Strict                 (Map, differenceWith, empty, filter, foldrWithKey, fromListWith,
                                                  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, sum, take, (!!), (/), (^))
import           Test.QuickCheck                 (Arbitrary (..))

import           ZkFold.Base.Algebra.Basic.Class
import           ZkFold.Base.Data.Vector         (Vector (..))
import           ZkFold.Prelude                  (elemIndex)

type Variable i = Ord i

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

-- | Monomial type
newtype M i j m = M m
    deriving ((forall x. M i j m -> Rep (M i j m) x)
-> (forall x. Rep (M i j m) x -> M i j m) -> Generic (M i j m)
forall x. Rep (M i j m) x -> M i j m
forall x. M i j m -> Rep (M i j m) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall k (i :: k) k (j :: k) m x. Rep (M i j m) x -> M i j m
forall k (i :: k) k (j :: k) m x. M i j m -> Rep (M i j m) x
$cfrom :: forall k (i :: k) k (j :: k) m x. M i j m -> Rep (M i j m) x
from :: forall x. M i j m -> Rep (M i j m) x
$cto :: forall k (i :: k) k (j :: k) m x. Rep (M i j m) x -> M i j m
to :: forall x. Rep (M i j m) x -> M i j m
Generic, M i j m -> ()
(M i j m -> ()) -> NFData (M i j m)
forall a. (a -> ()) -> NFData a
forall k (i :: k) k (j :: k) m. NFData m => M i j m -> ()
$crnf :: forall k (i :: k) k (j :: k) m. NFData m => M i j m -> ()
rnf :: M i j m -> ()
NFData, Maybe (M i j m)
Value -> Parser [M i j m]
Value -> Parser (M i j m)
(Value -> Parser (M i j m))
-> (Value -> Parser [M i j m])
-> Maybe (M i j m)
-> FromJSON (M i j m)
forall a.
(Value -> Parser a)
-> (Value -> Parser [a]) -> Maybe a -> FromJSON a
forall k (i :: k) k (j :: k) m. FromJSON m => Maybe (M i j m)
forall k (i :: k) k (j :: k) m.
FromJSON m =>
Value -> Parser [M i j m]
forall k (i :: k) k (j :: k) m.
FromJSON m =>
Value -> Parser (M i j m)
$cparseJSON :: forall k (i :: k) k (j :: k) m.
FromJSON m =>
Value -> Parser (M i j m)
parseJSON :: Value -> Parser (M i j m)
$cparseJSONList :: forall k (i :: k) k (j :: k) m.
FromJSON m =>
Value -> Parser [M i j m]
parseJSONList :: Value -> Parser [M i j m]
$comittedField :: forall k (i :: k) k (j :: k) m. FromJSON m => Maybe (M i j m)
omittedField :: Maybe (M i j m)
FromJSON, [M i j m] -> Value
[M i j m] -> Encoding
M i j m -> Bool
M i j m -> Value
M i j m -> Encoding
(M i j m -> Value)
-> (M i j m -> Encoding)
-> ([M i j m] -> Value)
-> ([M i j m] -> Encoding)
-> (M i j m -> Bool)
-> ToJSON (M i j m)
forall a.
(a -> Value)
-> (a -> Encoding)
-> ([a] -> Value)
-> ([a] -> Encoding)
-> (a -> Bool)
-> ToJSON a
forall k (i :: k) k (j :: k) m. ToJSON m => [M i j m] -> Value
forall k (i :: k) k (j :: k) m. ToJSON m => [M i j m] -> Encoding
forall k (i :: k) k (j :: k) m. ToJSON m => M i j m -> Bool
forall k (i :: k) k (j :: k) m. ToJSON m => M i j m -> Value
forall k (i :: k) k (j :: k) m. ToJSON m => M i j m -> Encoding
$ctoJSON :: forall k (i :: k) k (j :: k) m. ToJSON m => M i j m -> Value
toJSON :: M i j m -> Value
$ctoEncoding :: forall k (i :: k) k (j :: k) m. ToJSON m => M i j m -> Encoding
toEncoding :: M i j m -> Encoding
$ctoJSONList :: forall k (i :: k) k (j :: k) m. ToJSON m => [M i j m] -> Value
toJSONList :: [M i j m] -> Value
$ctoEncodingList :: forall k (i :: k) k (j :: k) m. ToJSON m => [M i j m] -> Encoding
toEncodingList :: [M i j m] -> Encoding
$comitField :: forall k (i :: k) k (j :: k) m. ToJSON m => M i j m -> Bool
omitField :: M i j m -> Bool
ToJSON)

-- | Most general type for a multivariate monomial
type Monomial' = M Natural Natural (Map Natural Natural)

type MonomialRepAny = Map Integer Integer

type MonomialRepBoundedDegree i d = Vector d (i, Bool)

type MonomialAny = M Integer Integer MonomialRepAny

type MonomialBoundedDegree i d = M i Bool (MonomialRepBoundedDegree i d)

instance Ord i => IsList (M i j (Map i j)) where
    type Item (M i j (Map i j)) = (i, j)
    toList :: M i j (Map i j) -> [Item (M i j (Map 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 (M i j (Map i j))] -> M i j (Map i j)
fromList [Item (M i j (Map i j))]
m = Map i j -> M i j (Map i j)
forall {k} {k} (i :: k) (j :: k) m. m -> M i j m
M (Map i j -> M i j (Map i j)) -> Map i j -> M i j (Map i j)
forall a b. (a -> b) -> a -> b
$ [Item (Map i j)] -> Map i j
forall l. IsList l => [Item l] -> l
fromList [Item (Map i j)]
[Item (M i j (Map i j))]
m

-- | Monomial constructor
monomial :: Monomial i j => Map i j -> M i j (Map i j)
monomial :: forall i j. Monomial i j => Map i j -> M i j (Map i j)
monomial = Map i j -> M i j (Map i j)
forall {k} {k} (i :: k) (j :: k) m. m -> M i j m
M (Map i j -> M i j (Map i j))
-> (Map i j -> Map i j) -> Map i j -> M i j (Map 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)

instance (Show i, Show j, Monomial i j) => Show (M i j (Map i j)) where
    show :: M i j (Map i j) -> String
show (M Map i j
m) = String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
"∙" ([String] -> String)
-> ([(i, j)] -> [String]) -> [(i, j)] -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((i, j) -> String) -> [(i, j)] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map (i, j) -> String
showVar ([(i, j)] -> String) -> [(i, j)] -> String
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) -> String
showVar (i
i, j
j) = String
"x" String -> ShowS
forall a. [a] -> [a] -> [a]
++ i -> String
forall a. Show a => a -> String
show i
i String -> 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 String
"" else String
"^" String -> ShowS
forall a. [a] -> [a] -> [a]
++ j -> String
forall a. Show a => a -> String
show j
j)

instance (Eq i, Eq j) => Eq (M i j (Map i j)) where
    M Map i j
asl == :: M i j (Map i j) -> M i j (Map 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 (Eq i, Ord i, Ord j) => Ord (M i j (Map i j)) where
    compare :: M i j (Map i j) -> M i j (Map 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 Arbitrary m => Arbitrary (M i j m) where
    arbitrary :: Gen (M i j m)
arbitrary = m -> M i j m
forall {k} {k} (i :: k) (j :: k) m. m -> M i j m
M (m -> M i j m) -> Gen m -> Gen (M i j m)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen m
forall a. Arbitrary a => Gen a
arbitrary

instance Monomial i j => MultiplicativeSemigroup (M i j (Map i j)) where
    M Map i j
l * :: M i j (Map i j) -> M i j (Map i j) -> M i j (Map i j)
* M Map i j
r = Map i j -> M i j (Map i j)
forall {k} {k} (i :: k) (j :: k) m. m -> M i j m
M (Map i j -> M i j (Map i j)) -> Map i j -> M i j (Map 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 (M i j (Map i j)) Natural where
    ^ :: M i j (Map i j) -> Natural -> M i j (Map i j)
(^) = M i j (Map i j) -> Natural -> M i j (Map i j)
forall a. MultiplicativeMonoid a => a -> Natural -> a
natPow

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

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

instance (Monomial i j, Ring j) => MultiplicativeGroup (M i j (Map i j)) where
    invert :: M i j (Map i j) -> M i j (Map i j)
invert (M Map i j
m) = Map i j -> M i j (Map i j)
forall {k} {k} (i :: k) (j :: k) m. m -> M i j m
M (Map i j -> M i j (Map i j)) -> Map i j -> M i j (Map 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 / :: M i j (Map i j) -> M i j (Map i j) -> M i j (Map i j)
/ M Map i j
r = Map i j -> M i j (Map i j)
forall {k} {k} (i :: k) (j :: k) m. m -> M i j m
M (Map i j -> M i j (Map i j)) -> Map i j -> M i j (Map 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)

evalMapM :: forall i j b .
    MultiplicativeMonoid b =>
    Exponent b j =>
    (i -> b) -> M i j (Map i j) -> b
evalMapM :: forall i j b.
(MultiplicativeMonoid b, Exponent b j) =>
(i -> b) -> M i j (Map i j) -> b
evalMapM 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

evalVectorM :: forall i j b d .
    Monomial i j =>
    MultiplicativeMonoid b =>
    Exponent b j =>
    (i -> b) -> M i j (Vector d (i, Bool)) -> b
evalVectorM :: forall i j b (d :: Natural).
(Monomial i j, MultiplicativeMonoid b, Exponent b j) =>
(i -> b) -> M i j (Vector d (i, Bool)) -> b
evalVectorM i -> b
f (M (Vector [(i, Bool)]
v)) =
    (i -> b) -> M i j (Map i j) -> b
forall i j b.
(MultiplicativeMonoid b, Exponent b j) =>
(i -> b) -> M i j (Map i j) -> b
evalMapM i -> b
f (M i j (Map i j) -> b)
-> ([(i, j)] -> M i j (Map i j)) -> [(i, j)] -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map i j -> M i j (Map i j)
forall {k} {k} (i :: k) (j :: k) m. m -> M i j m
M (Map i j -> M i j (Map i j))
-> ([(i, j)] -> Map i j) -> [(i, j)] -> M i j (Map i j)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (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)] -> b) -> [(i, j)] -> b
forall a b. (a -> b) -> a -> b
$ ((i, Bool) -> [(i, j)] -> [(i, j)])
-> [(i, j)] -> [(i, Bool)] -> [(i, j)]
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: Type -> Type) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\(i
i, Bool
x) [(i, j)]
xs -> if Bool
x then (i
i, forall a. MultiplicativeMonoid a => a
one @j) (i, j) -> [(i, j)] -> [(i, j)]
forall a. a -> [a] -> [a]
: [(i, j)]
xs else [(i, j)]
xs) [] [(i, Bool)]
v

mapVar :: [Natural] -> Natural -> Natural
mapVar :: [Natural] -> Natural -> Natural
mapVar [Natural]
vars Natural
x = case Natural
x Natural -> [Natural] -> Maybe Natural
forall a. Eq a => a -> [a] -> Maybe Natural
`elemIndex` [Natural]
vars of
    Just Natural
i  -> Natural
i
    Maybe Natural
Nothing -> String -> Natural
forall a. HasCallStack => String -> a
error String
"mapVar: something went wrong"

mapVarMonomial :: [Natural] -> Monomial' -> Monomial'
mapVarMonomial :: [Natural] -> Monomial' -> Monomial'
mapVarMonomial [Natural]
vars (M Map Natural Natural
as) = Map Natural Natural -> Monomial'
forall {k} {k} (i :: k) (j :: k) m. m -> M i j m
M (Map Natural Natural -> Monomial')
-> Map Natural Natural -> Monomial'
forall a b. (a -> b) -> a -> b
$ (Natural -> Natural) -> Map Natural Natural -> Map Natural Natural
forall k2 k1 a. Ord k2 => (k1 -> k2) -> Map k1 a -> Map k2 a
mapKeys ([Natural] -> Natural -> Natural
mapVar [Natural]
vars) Map Natural Natural
as