module ZkFold.Base.Data.Sparse.Matrix where

import           Data.Map                         (Map)
import           Data.Zip                         (Semialign (..), Zip (..))
import           Prelude                          hiding (length, sum, zip, zipWith, (*))
import           Test.QuickCheck                  (Arbitrary (..))

import           ZkFold.Base.Algebra.Basic.Field  (Zp)
import           ZkFold.Base.Algebra.Basic.Number (KnownNat)

newtype SMatrix m n a = SMatrix { forall (m :: Natural) (n :: Natural) a.
SMatrix m n a -> Map (Zp m, Zp n) a
fromSMatrix :: Map (Zp m, Zp n) a }
    deriving (Int -> SMatrix m n a -> ShowS
[SMatrix m n a] -> ShowS
SMatrix m n a -> String
(Int -> SMatrix m n a -> ShowS)
-> (SMatrix m n a -> String)
-> ([SMatrix m n a] -> ShowS)
-> Show (SMatrix m n a)
forall (m :: Natural) (n :: Natural) a.
Show a =>
Int -> SMatrix m n a -> ShowS
forall (m :: Natural) (n :: Natural) a.
Show a =>
[SMatrix m n a] -> ShowS
forall (m :: Natural) (n :: Natural) a.
Show a =>
SMatrix m n a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall (m :: Natural) (n :: Natural) a.
Show a =>
Int -> SMatrix m n a -> ShowS
showsPrec :: Int -> SMatrix m n a -> ShowS
$cshow :: forall (m :: Natural) (n :: Natural) a.
Show a =>
SMatrix m n a -> String
show :: SMatrix m n a -> String
$cshowList :: forall (m :: Natural) (n :: Natural) a.
Show a =>
[SMatrix m n a] -> ShowS
showList :: [SMatrix m n a] -> ShowS
Show, SMatrix m n a -> SMatrix m n a -> Bool
(SMatrix m n a -> SMatrix m n a -> Bool)
-> (SMatrix m n a -> SMatrix m n a -> Bool) -> Eq (SMatrix m n a)
forall (m :: Natural) (n :: Natural) a.
(KnownNat m, KnownNat n, Eq a) =>
SMatrix m n a -> SMatrix m n a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall (m :: Natural) (n :: Natural) a.
(KnownNat m, KnownNat n, Eq a) =>
SMatrix m n a -> SMatrix m n a -> Bool
== :: SMatrix m n a -> SMatrix m n a -> Bool
$c/= :: forall (m :: Natural) (n :: Natural) a.
(KnownNat m, KnownNat n, Eq a) =>
SMatrix m n a -> SMatrix m n a -> Bool
/= :: SMatrix m n a -> SMatrix m n a -> Bool
Eq)

instance Foldable (SMatrix m n) where
    foldr :: forall a b. (a -> b -> b) -> b -> SMatrix m n a -> b
foldr a -> b -> b
f b
z (SMatrix Map (Zp m, Zp n) a
as) = (a -> b -> b) -> b -> Map (Zp m, Zp n) a -> b
forall a b. (a -> b -> b) -> b -> Map (Zp m, Zp n) a -> b
forall (t :: Type -> Type) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> b -> b
f b
z Map (Zp m, Zp n) a
as

instance Functor (SMatrix m n) where
    fmap :: forall a b. (a -> b) -> SMatrix m n a -> SMatrix m n b
fmap a -> b
f (SMatrix Map (Zp m, Zp n) a
as) = Map (Zp m, Zp n) b -> SMatrix m n b
forall (m :: Natural) (n :: Natural) a.
Map (Zp m, Zp n) a -> SMatrix m n a
SMatrix (Map (Zp m, Zp n) b -> SMatrix m n b)
-> Map (Zp m, Zp n) b -> SMatrix m n b
forall a b. (a -> b) -> a -> b
$ (a -> b) -> Map (Zp m, Zp n) a -> Map (Zp m, Zp n) b
forall a b. (a -> b) -> Map (Zp m, Zp n) a -> Map (Zp m, Zp n) b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f Map (Zp m, Zp n) a
as

instance (KnownNat m, KnownNat n) => Semialign (SMatrix m n) where
    align :: forall a b.
SMatrix m n a -> SMatrix m n b -> SMatrix m n (These a b)
align (SMatrix Map (Zp m, Zp n) a
as) (SMatrix Map (Zp m, Zp n) b
bs) = Map (Zp m, Zp n) (These a b) -> SMatrix m n (These a b)
forall (m :: Natural) (n :: Natural) a.
Map (Zp m, Zp n) a -> SMatrix m n a
SMatrix (Map (Zp m, Zp n) (These a b) -> SMatrix m n (These a b))
-> Map (Zp m, Zp n) (These a b) -> SMatrix m n (These a b)
forall a b. (a -> b) -> a -> b
$ Map (Zp m, Zp n) a
-> Map (Zp m, Zp n) b -> Map (Zp m, Zp n) (These a b)
forall a b.
Map (Zp m, Zp n) a
-> Map (Zp m, Zp n) b -> Map (Zp m, Zp n) (These a b)
forall (f :: Type -> Type) a b.
Semialign f =>
f a -> f b -> f (These a b)
align Map (Zp m, Zp n) a
as Map (Zp m, Zp n) b
bs

    alignWith :: forall a b c.
(These a b -> c) -> SMatrix m n a -> SMatrix m n b -> SMatrix m n c
alignWith These a b -> c
f (SMatrix Map (Zp m, Zp n) a
as) (SMatrix Map (Zp m, Zp n) b
bs) = Map (Zp m, Zp n) c -> SMatrix m n c
forall (m :: Natural) (n :: Natural) a.
Map (Zp m, Zp n) a -> SMatrix m n a
SMatrix (Map (Zp m, Zp n) c -> SMatrix m n c)
-> Map (Zp m, Zp n) c -> SMatrix m n c
forall a b. (a -> b) -> a -> b
$ (These a b -> c)
-> Map (Zp m, Zp n) a -> Map (Zp m, Zp n) b -> Map (Zp m, Zp n) c
forall a b c.
(These a b -> c)
-> Map (Zp m, Zp n) a -> Map (Zp m, Zp n) b -> Map (Zp m, Zp n) c
forall (f :: Type -> Type) a b c.
Semialign f =>
(These a b -> c) -> f a -> f b -> f c
alignWith These a b -> c
f Map (Zp m, Zp n) a
as Map (Zp m, Zp n) b
bs

instance (KnownNat m, KnownNat n) => Zip (SMatrix m n) where
    zip :: forall a b. SMatrix m n a -> SMatrix m n b -> SMatrix m n (a, b)
zip (SMatrix Map (Zp m, Zp n) a
as) (SMatrix Map (Zp m, Zp n) b
bs) = Map (Zp m, Zp n) (a, b) -> SMatrix m n (a, b)
forall (m :: Natural) (n :: Natural) a.
Map (Zp m, Zp n) a -> SMatrix m n a
SMatrix (Map (Zp m, Zp n) (a, b) -> SMatrix m n (a, b))
-> Map (Zp m, Zp n) (a, b) -> SMatrix m n (a, b)
forall a b. (a -> b) -> a -> b
$ Map (Zp m, Zp n) a -> Map (Zp m, Zp n) b -> Map (Zp m, Zp n) (a, b)
forall a b.
Map (Zp m, Zp n) a -> Map (Zp m, Zp n) b -> Map (Zp m, Zp n) (a, b)
forall (f :: Type -> Type) a b. Zip f => f a -> f b -> f (a, b)
zip Map (Zp m, Zp n) a
as Map (Zp m, Zp n) b
bs

    zipWith :: forall a b c.
(a -> b -> c) -> SMatrix m n a -> SMatrix m n b -> SMatrix m n c
zipWith a -> b -> c
f (SMatrix Map (Zp m, Zp n) a
as) (SMatrix Map (Zp m, Zp n) b
bs) = Map (Zp m, Zp n) c -> SMatrix m n c
forall (m :: Natural) (n :: Natural) a.
Map (Zp m, Zp n) a -> SMatrix m n a
SMatrix (Map (Zp m, Zp n) c -> SMatrix m n c)
-> Map (Zp m, Zp n) c -> SMatrix m n c
forall a b. (a -> b) -> a -> b
$ (a -> b -> c)
-> Map (Zp m, Zp n) a -> Map (Zp m, Zp n) b -> Map (Zp m, Zp n) c
forall a b c.
(a -> b -> c)
-> Map (Zp m, Zp n) a -> Map (Zp m, Zp n) b -> Map (Zp m, Zp n) c
forall (f :: Type -> Type) a b c.
Zip f =>
(a -> b -> c) -> f a -> f b -> f c
zipWith a -> b -> c
f Map (Zp m, Zp n) a
as Map (Zp m, Zp n) b
bs

instance (KnownNat m, KnownNat n, Arbitrary a) => Arbitrary (SMatrix m n a) where
    arbitrary :: Gen (SMatrix m n a)
arbitrary = Map (Zp m, Zp n) a -> SMatrix m n a
forall (m :: Natural) (n :: Natural) a.
Map (Zp m, Zp n) a -> SMatrix m n a
SMatrix (Map (Zp m, Zp n) a -> SMatrix m n a)
-> Gen (Map (Zp m, Zp n) a) -> Gen (SMatrix m n a)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen (Map (Zp m, Zp n) a)
forall a. Arbitrary a => Gen a
arbitrary