{-# LANGUAGE TypeApplications #-}

module ZkFold.Base.Data.Matrix where

import           Data.Bifunctor                   (first)
import qualified Data.List                        as List
import           Data.Maybe                       (fromJust)
import           Data.These
import           Data.Zip                         (Semialign (..), Zip (..))
import           Prelude                          hiding (Num (..), length, sum, zip, zipWith)
import           System.Random                    (Random (..))
import           Test.QuickCheck                  (Arbitrary (..))

import           ZkFold.Base.Algebra.Basic.Class
import           ZkFold.Base.Algebra.Basic.Number
import           ZkFold.Base.Data.Vector

-- TODO: implement a proper matrix algebra
-- Could be useful for speeding up the proof computations

newtype Matrix m n a = Matrix (Vector m (Vector n a))
    deriving (Int -> Matrix m n a -> ShowS
[Matrix m n a] -> ShowS
Matrix m n a -> String
(Int -> Matrix m n a -> ShowS)
-> (Matrix m n a -> String)
-> ([Matrix m n a] -> ShowS)
-> Show (Matrix m n a)
forall (m :: Natural) (n :: Natural) a.
Show a =>
Int -> Matrix m n a -> ShowS
forall (m :: Natural) (n :: Natural) a.
Show a =>
[Matrix m n a] -> ShowS
forall (m :: Natural) (n :: Natural) a.
Show a =>
Matrix 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 -> Matrix m n a -> ShowS
showsPrec :: Int -> Matrix m n a -> ShowS
$cshow :: forall (m :: Natural) (n :: Natural) a.
Show a =>
Matrix m n a -> String
show :: Matrix m n a -> String
$cshowList :: forall (m :: Natural) (n :: Natural) a.
Show a =>
[Matrix m n a] -> ShowS
showList :: [Matrix m n a] -> ShowS
Show, Matrix m n a -> Matrix m n a -> Bool
(Matrix m n a -> Matrix m n a -> Bool)
-> (Matrix m n a -> Matrix m n a -> Bool) -> Eq (Matrix m n a)
forall (m :: Natural) (n :: Natural) a.
Eq a =>
Matrix m n a -> Matrix m n a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall (m :: Natural) (n :: Natural) a.
Eq a =>
Matrix m n a -> Matrix m n a -> Bool
== :: Matrix m n a -> Matrix m n a -> Bool
$c/= :: forall (m :: Natural) (n :: Natural) a.
Eq a =>
Matrix m n a -> Matrix m n a -> Bool
/= :: Matrix m n a -> Matrix m n a -> Bool
Eq)

toMatrix :: forall m n a . (KnownNat m, KnownNat n) => [[a]] -> Maybe (Matrix m n a)
toMatrix :: forall (m :: Natural) (n :: Natural) a.
(KnownNat m, KnownNat n) =>
[[a]] -> Maybe (Matrix m n a)
toMatrix [[a]]
as = do
    [Vector n a]
as' <- ([a] -> Maybe (Vector n a)) -> [[a]] -> Maybe [Vector n a]
forall (t :: Type -> Type) (m :: Type -> Type) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: Type -> Type) a b.
Monad m =>
(a -> m b) -> [a] -> m [b]
mapM (forall (size :: Natural) a.
KnownNat size =>
[a] -> Maybe (Vector size a)
toVector @n) [[a]]
as
    Vector m (Vector n a) -> Matrix m n a
forall (m :: Natural) (n :: Natural) a.
Vector m (Vector n a) -> Matrix m n a
Matrix (Vector m (Vector n a) -> Matrix m n a)
-> Maybe (Vector m (Vector n a)) -> Maybe (Matrix m n a)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (size :: Natural) a.
KnownNat size =>
[a] -> Maybe (Vector size a)
toVector @m [Vector n a]
as'

fromMatrix :: forall m n a . Matrix m n a -> [[a]]
fromMatrix :: forall (m :: Natural) (n :: Natural) a. Matrix m n a -> [[a]]
fromMatrix (Matrix Vector m (Vector n a)
as) = (Vector n a -> [a]) -> [Vector n a] -> [[a]]
forall a b. (a -> b) -> [a] -> [b]
map Vector n a -> [a]
forall (size :: Natural) a. Vector size a -> [a]
fromVector ([Vector n a] -> [[a]]) -> [Vector n a] -> [[a]]
forall a b. (a -> b) -> a -> b
$ Vector m (Vector n a) -> [Vector n a]
forall (size :: Natural) a. Vector size a -> [a]
fromVector Vector m (Vector n a)
as

transpose :: forall m n a . (KnownNat m, KnownNat n) => Matrix m n a -> Matrix n m a
transpose :: forall (m :: Natural) (n :: Natural) a.
(KnownNat m, KnownNat n) =>
Matrix m n a -> Matrix n m a
transpose Matrix m n a
m = Maybe (Matrix n m a) -> Matrix n m a
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe (Matrix n m a) -> Matrix n m a)
-> Maybe (Matrix n m a) -> Matrix n m a
forall a b. (a -> b) -> a -> b
$ forall (m :: Natural) (n :: Natural) a.
(KnownNat m, KnownNat n) =>
[[a]] -> Maybe (Matrix m n a)
toMatrix @n @m ([[a]] -> Maybe (Matrix n m a)) -> [[a]] -> Maybe (Matrix n m a)
forall a b. (a -> b) -> a -> b
$ [[a]] -> [[a]]
forall a. [[a]] -> [[a]]
List.transpose ([[a]] -> [[a]]) -> [[a]] -> [[a]]
forall a b. (a -> b) -> a -> b
$ Matrix m n a -> [[a]]
forall (m :: Natural) (n :: Natural) a. Matrix m n a -> [[a]]
fromMatrix Matrix m n a
m

outer :: forall m n a b c. (a -> b -> c) -> Vector m a -> Vector n b -> Matrix m n c
outer :: forall (m :: Natural) (n :: Natural) a b c.
(a -> b -> c) -> Vector m a -> Vector n b -> Matrix m n c
outer a -> b -> c
f Vector m a
a Vector n b
b = Vector m (Vector n c) -> Matrix m n c
forall (m :: Natural) (n :: Natural) a.
Vector m (Vector n a) -> Matrix m n a
Matrix (Vector m (Vector n c) -> Matrix m n c)
-> Vector m (Vector n c) -> Matrix m n c
forall a b. (a -> b) -> a -> b
$ (a -> Vector n c) -> Vector m a -> Vector m (Vector n c)
forall a b. (a -> b) -> Vector m a -> Vector m b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap (\a
x -> (b -> c) -> Vector n b -> Vector n c
forall a b. (a -> b) -> Vector n a -> Vector n b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap (a -> b -> c
f a
x) Vector n b
b) Vector m a
a

-- | Hadamard (entry-wise) matrix product
(.*) :: MultiplicativeSemigroup a => Matrix m n a -> Matrix m n a -> Matrix m n a
.* :: forall a (m :: Natural) (n :: Natural).
MultiplicativeSemigroup a =>
Matrix m n a -> Matrix m n a -> Matrix m n a
(.*) = (a -> a -> a) -> Matrix m n a -> Matrix m n a -> Matrix m n a
forall a b c.
(a -> b -> c) -> Matrix m n a -> Matrix m n b -> Matrix m n c
forall (f :: Type -> Type) a b c.
Zip f =>
(a -> b -> c) -> f a -> f b -> f c
zipWith a -> a -> a
forall a. MultiplicativeSemigroup a => a -> a -> a
(*)

sum1 :: (Semiring a) => Matrix m n a -> Vector n a
sum1 :: forall a (m :: Natural) (n :: Natural).
Semiring a =>
Matrix m n a -> Vector n a
sum1 (Matrix Vector m (Vector n a)
as) = [a] -> Vector n a
forall (size :: Natural) a. [a] -> Vector size a
Vector (Vector n a -> a
forall (t :: Type -> Type) a.
(Foldable t, AdditiveMonoid a) =>
t a -> a
sum (Vector n a -> a) -> [Vector n a] -> [a]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Vector m (Vector n a) -> [Vector n a]
forall (size :: Natural) a. Vector size a -> [a]
fromVector Vector m (Vector n a)
as)

sum2 :: (KnownNat m, KnownNat n, Semiring a) => Matrix m n a -> Vector m a
sum2 :: forall (m :: Natural) (n :: Natural) a.
(KnownNat m, KnownNat n, Semiring a) =>
Matrix m n a -> Vector m a
sum2 (Matrix Vector m (Vector n a)
as) = Matrix n m a -> Vector m a
forall a (m :: Natural) (n :: Natural).
Semiring a =>
Matrix m n a -> Vector n a
sum1 (Matrix n m a -> Vector m a) -> Matrix n m a -> Vector m a
forall a b. (a -> b) -> a -> b
$ Matrix m n a -> Matrix n m a
forall (m :: Natural) (n :: Natural) a.
(KnownNat m, KnownNat n) =>
Matrix m n a -> Matrix n m a
transpose (Matrix m n a -> Matrix n m a) -> Matrix m n a -> Matrix n m a
forall a b. (a -> b) -> a -> b
$ Vector m (Vector n a) -> Matrix m n a
forall (m :: Natural) (n :: Natural) a.
Vector m (Vector n a) -> Matrix m n a
Matrix Vector m (Vector n a)
as

matrixDotProduct :: forall m n a . Semiring a => Matrix m n a -> Matrix m n a -> a
matrixDotProduct :: forall (m :: Natural) (n :: Natural) a.
Semiring a =>
Matrix m n a -> Matrix m n a -> a
matrixDotProduct Matrix m n a
a Matrix m n a
b = let Matrix Vector m (Vector n a)
m = Matrix m n a
a Matrix m n a -> Matrix m n a -> Matrix m n a
forall a (m :: Natural) (n :: Natural).
MultiplicativeSemigroup a =>
Matrix m n a -> Matrix m n a -> Matrix m n a
.* Matrix m n a
b in Vector m a -> a
forall (t :: Type -> Type) a.
(Foldable t, AdditiveMonoid a) =>
t a -> a
sum (Vector m a -> a) -> Vector m a -> a
forall a b. (a -> b) -> a -> b
$ (Vector n a -> a) -> Vector m (Vector n a) -> Vector m a
forall a b. (a -> b) -> Vector m a -> Vector m b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap Vector n a -> a
forall (t :: Type -> Type) a.
(Foldable t, AdditiveMonoid a) =>
t a -> a
sum Vector m (Vector n a)
m

-- -- | Matrix multiplication
(.*.) :: (KnownNat n, KnownNat k, Semiring a) => Matrix m n a -> Matrix n k a -> Matrix m k a
Matrix m n a
a .*. :: forall (n :: Natural) (k :: Natural) a (m :: Natural).
(KnownNat n, KnownNat k, Semiring a) =>
Matrix m n a -> Matrix n k a -> Matrix m k a
.*. Matrix n k a
b =
    let Matrix Vector m (Vector n a)
a' = Matrix m n a
a
        Matrix Vector k (Vector n a)
b' = Matrix n k a -> Matrix k n a
forall (m :: Natural) (n :: Natural) a.
(KnownNat m, KnownNat n) =>
Matrix m n a -> Matrix n m a
transpose Matrix n k a
b
    in Vector m (Vector k a) -> Matrix m k a
forall (m :: Natural) (n :: Natural) a.
Vector m (Vector n a) -> Matrix m n a
Matrix (Vector m (Vector k a) -> Matrix m k a)
-> Vector m (Vector k a) -> Matrix m k a
forall a b. (a -> b) -> a -> b
$ (Vector n a -> Vector k a)
-> Vector m (Vector n a) -> Vector m (Vector k a)
forall a b. (a -> b) -> Vector m a -> Vector m b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap (\Vector n a
x -> (Vector n a -> a) -> Vector k (Vector n a) -> Vector k a
forall a b. (a -> b) -> Vector k a -> Vector k b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap (Vector n a -> Vector n a -> a
forall (size :: Natural) a.
Semiring a =>
Vector size a -> Vector size a -> a
vectorDotProduct Vector n a
x) Vector k (Vector n a)
b') Vector m (Vector n a)
a'

instance Functor (Matrix m n) where
    fmap :: forall a b. (a -> b) -> Matrix m n a -> Matrix m n b
fmap a -> b
f (Matrix Vector m (Vector n a)
as) = Vector m (Vector n b) -> Matrix m n b
forall (m :: Natural) (n :: Natural) a.
Vector m (Vector n a) -> Matrix m n a
Matrix (Vector m (Vector n b) -> Matrix m n b)
-> Vector m (Vector n b) -> Matrix m n b
forall a b. (a -> b) -> a -> b
$ (Vector n a -> Vector n b)
-> Vector m (Vector n a) -> Vector m (Vector n b)
forall a b. (a -> b) -> Vector m a -> Vector m b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap ((a -> b) -> Vector n a -> Vector n b
forall a b. (a -> b) -> Vector n a -> Vector n b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f) Vector m (Vector n a)
as

instance (KnownNat m, KnownNat n) => Applicative (Matrix m n) where
    pure :: forall a. a -> Matrix m n a
pure a
a = Vector m (Vector n a) -> Matrix m n a
forall (m :: Natural) (n :: Natural) a.
Vector m (Vector n a) -> Matrix m n a
Matrix (Vector m (Vector n a) -> Matrix m n a)
-> Vector m (Vector n a) -> Matrix m n a
forall a b. (a -> b) -> a -> b
$ Vector n a -> Vector m (Vector n a)
forall a. a -> Vector m a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Vector n a -> Vector m (Vector n a))
-> Vector n a -> Vector m (Vector n a)
forall a b. (a -> b) -> a -> b
$ a -> Vector n a
forall a. a -> Vector n a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure a
a

    (Matrix Vector m (Vector n (a -> b))
fs) <*> :: forall a b. Matrix m n (a -> b) -> Matrix m n a -> Matrix m n b
<*> (Matrix Vector m (Vector n a)
as) = Vector m (Vector n b) -> Matrix m n b
forall (m :: Natural) (n :: Natural) a.
Vector m (Vector n a) -> Matrix m n a
Matrix (Vector m (Vector n b) -> Matrix m n b)
-> Vector m (Vector n b) -> Matrix m n b
forall a b. (a -> b) -> a -> b
$ (Vector n (a -> b) -> Vector n a -> Vector n b)
-> Vector m (Vector n (a -> b))
-> Vector m (Vector n a)
-> Vector m (Vector n b)
forall a b c.
(a -> b -> c) -> Vector m a -> Vector m b -> Vector m c
forall (f :: Type -> Type) a b c.
Zip f =>
(a -> b -> c) -> f a -> f b -> f c
zipWith Vector n (a -> b) -> Vector n a -> Vector n b
forall a b. Vector n (a -> b) -> Vector n a -> Vector n b
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
(<*>) Vector m (Vector n (a -> b))
fs Vector m (Vector n a)
as

instance Semialign (Matrix m n) where
    align :: forall a b. Matrix m n a -> Matrix m n b -> Matrix m n (These a b)
align (Matrix Vector m (Vector n a)
as) (Matrix Vector m (Vector n b)
bs) = Vector m (Vector n (These a b)) -> Matrix m n (These a b)
forall (m :: Natural) (n :: Natural) a.
Vector m (Vector n a) -> Matrix m n a
Matrix (Vector m (Vector n (These a b)) -> Matrix m n (These a b))
-> Vector m (Vector n (These a b)) -> Matrix m n (These a b)
forall a b. (a -> b) -> a -> b
$ (Vector n a -> Vector n b -> Vector n (These a b))
-> Vector m (Vector n a)
-> Vector m (Vector n b)
-> Vector m (Vector n (These a b))
forall a b c.
(a -> b -> c) -> Vector m a -> Vector m b -> Vector m c
forall (f :: Type -> Type) a b c.
Zip f =>
(a -> b -> c) -> f a -> f b -> f c
zipWith ((a -> b -> These a b)
-> Vector n a -> Vector n b -> Vector n (These a b)
forall a b c.
(a -> b -> c) -> Vector n a -> Vector n b -> Vector n c
forall (f :: Type -> Type) a b c.
Zip f =>
(a -> b -> c) -> f a -> f b -> f c
zipWith a -> b -> These a b
forall a b. a -> b -> These a b
These) Vector m (Vector n a)
as Vector m (Vector n b)
bs

    alignWith :: forall a b c.
(These a b -> c) -> Matrix m n a -> Matrix m n b -> Matrix m n c
alignWith These a b -> c
f (Matrix Vector m (Vector n a)
as) (Matrix Vector m (Vector n b)
bs) = Vector m (Vector n c) -> Matrix m n c
forall (m :: Natural) (n :: Natural) a.
Vector m (Vector n a) -> Matrix m n a
Matrix (Vector m (Vector n c) -> Matrix m n c)
-> Vector m (Vector n c) -> Matrix m n c
forall a b. (a -> b) -> a -> b
$ (Vector n a -> Vector n b -> Vector n c)
-> Vector m (Vector n a)
-> Vector m (Vector n b)
-> Vector m (Vector n c)
forall a b c.
(a -> b -> c) -> Vector m a -> Vector m b -> Vector m c
forall (f :: Type -> Type) a b c.
Zip f =>
(a -> b -> c) -> f a -> f b -> f c
zipWith ((a -> b -> c) -> Vector n a -> Vector n b -> Vector n c
forall a b c.
(a -> b -> c) -> Vector n a -> Vector n b -> Vector n c
forall (f :: Type -> Type) a b c.
Zip f =>
(a -> b -> c) -> f a -> f b -> f c
zipWith (\a
a b
b -> These a b -> c
f (These a b -> c) -> These a b -> c
forall a b. (a -> b) -> a -> b
$ a -> b -> These a b
forall a b. a -> b -> These a b
These a
a b
b)) Vector m (Vector n a)
as Vector m (Vector n b)
bs

instance Zip (Matrix m n) where
    zip :: forall a b. Matrix m n a -> Matrix m n b -> Matrix m n (a, b)
zip (Matrix Vector m (Vector n a)
as) (Matrix Vector m (Vector n b)
bs) = Vector m (Vector n (a, b)) -> Matrix m n (a, b)
forall (m :: Natural) (n :: Natural) a.
Vector m (Vector n a) -> Matrix m n a
Matrix (Vector m (Vector n (a, b)) -> Matrix m n (a, b))
-> Vector m (Vector n (a, b)) -> Matrix m n (a, b)
forall a b. (a -> b) -> a -> b
$ (Vector n a -> Vector n b -> Vector n (a, b))
-> Vector m (Vector n a)
-> Vector m (Vector n b)
-> Vector m (Vector n (a, b))
forall a b c.
(a -> b -> c) -> Vector m a -> Vector m b -> Vector m c
forall (f :: Type -> Type) a b c.
Zip f =>
(a -> b -> c) -> f a -> f b -> f c
zipWith Vector n a -> Vector n b -> Vector n (a, b)
forall a b. Vector n a -> Vector n b -> Vector n (a, b)
forall (f :: Type -> Type) a b. Zip f => f a -> f b -> f (a, b)
zip Vector m (Vector n a)
as Vector m (Vector n b)
bs

    zipWith :: forall a b c.
(a -> b -> c) -> Matrix m n a -> Matrix m n b -> Matrix m n c
zipWith a -> b -> c
f (Matrix Vector m (Vector n a)
as) (Matrix Vector m (Vector n b)
bs) = Vector m (Vector n c) -> Matrix m n c
forall (m :: Natural) (n :: Natural) a.
Vector m (Vector n a) -> Matrix m n a
Matrix (Vector m (Vector n c) -> Matrix m n c)
-> Vector m (Vector n c) -> Matrix m n c
forall a b. (a -> b) -> a -> b
$ (Vector n a -> Vector n b -> Vector n c)
-> Vector m (Vector n a)
-> Vector m (Vector n b)
-> Vector m (Vector n c)
forall a b c.
(a -> b -> c) -> Vector m a -> Vector m b -> Vector m c
forall (f :: Type -> Type) a b c.
Zip f =>
(a -> b -> c) -> f a -> f b -> f c
zipWith ((a -> b -> c) -> Vector n a -> Vector n b -> Vector n c
forall a b c.
(a -> b -> c) -> Vector n a -> Vector n b -> Vector 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) Vector m (Vector n a)
as Vector m (Vector n b)
bs

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

instance (Random a, KnownNat m, KnownNat n) => Random (Matrix m n a) where
    random :: forall g. RandomGen g => g -> (Matrix m n a, g)
random g
g =
        let as :: ([Vector n a], g)
as = (([Vector n a], g) -> Natural -> ([Vector n a], g))
-> ([Vector n a], g) -> [Natural] -> ([Vector n a], g)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: Type -> Type) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\([Vector n a]
as', g
g') Natural
_ ->
                let (Vector n a
a, g
g'') = g -> (Vector n a, g)
forall g. RandomGen g => g -> (Vector n a, g)
forall a g. (Random a, RandomGen g) => g -> (a, g)
random g
g'
                in ([Vector n a]
as' [Vector n a] -> [Vector n a] -> [Vector n a]
forall a. [a] -> [a] -> [a]
++ [Vector n a
a], g
g''))
                ([], g
g) [Natural
1..forall (n :: Natural). KnownNat n => Natural
value @m]
        in ([Vector n a] -> Matrix m n a)
-> ([Vector n a], g) -> (Matrix m n a, g)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: Type -> Type -> Type) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (Vector m (Vector n a) -> Matrix m n a
forall (m :: Natural) (n :: Natural) a.
Vector m (Vector n a) -> Matrix m n a
Matrix (Vector m (Vector n a) -> Matrix m n a)
-> ([Vector n a] -> Vector m (Vector n a))
-> [Vector n a]
-> Matrix m n a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Vector n a] -> Vector m (Vector n a)
forall (size :: Natural) a. [a] -> Vector size a
Vector) ([Vector n a], g)
as

    randomR :: forall g.
RandomGen g =>
(Matrix m n a, Matrix m n a) -> g -> (Matrix m n a, g)
randomR (Matrix Vector m (Vector n a)
xs, Matrix Vector m (Vector n a)
ys) g
g =
        let as :: ([Vector n a], g)
as = (([Vector n a], g), ([Vector n a], [Vector n a]))
-> ([Vector n a], g)
forall a b. (a, b) -> a
fst ((([Vector n a], g), ([Vector n a], [Vector n a]))
 -> ([Vector n a], g))
-> (([Vector n a], g), ([Vector n a], [Vector n a]))
-> ([Vector n a], g)
forall a b. (a -> b) -> a -> b
$ ((([Vector n a], g), ([Vector n a], [Vector n a]))
 -> Natural -> (([Vector n a], g), ([Vector n a], [Vector n a])))
-> (([Vector n a], g), ([Vector n a], [Vector n a]))
-> [Natural]
-> (([Vector n a], g), ([Vector n a], [Vector n a]))
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: Type -> Type) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\(([Vector n a]
as', g
g'), ([Vector n a]
xs', [Vector n a]
ys')) Natural
_ ->
                let (Vector n a
a, g
g'') = (Vector n a, Vector n a) -> g -> (Vector n a, g)
forall g.
RandomGen g =>
(Vector n a, Vector n a) -> g -> (Vector n a, g)
forall a g. (Random a, RandomGen g) => (a, a) -> g -> (a, g)
randomR ([Vector n a] -> Vector n a
forall a. HasCallStack => [a] -> a
head [Vector n a]
xs', [Vector n a] -> Vector n a
forall a. HasCallStack => [a] -> a
head [Vector n a]
ys') g
g'
                in (([Vector n a]
as' [Vector n a] -> [Vector n a] -> [Vector n a]
forall a. [a] -> [a] -> [a]
++ [Vector n a
a], g
g''), ([Vector n a] -> [Vector n a]
forall a. HasCallStack => [a] -> [a]
tail [Vector n a]
xs', [Vector n a] -> [Vector n a]
forall a. HasCallStack => [a] -> [a]
tail [Vector n a]
ys')))
                (([], g
g), (Vector m (Vector n a) -> [Vector n a]
forall (size :: Natural) a. Vector size a -> [a]
fromVector Vector m (Vector n a)
xs, Vector m (Vector n a) -> [Vector n a]
forall (size :: Natural) a. Vector size a -> [a]
fromVector Vector m (Vector n a)
ys)) [Natural
1..forall (n :: Natural). KnownNat n => Natural
value @m]
        in ([Vector n a] -> Matrix m n a)
-> ([Vector n a], g) -> (Matrix m n a, g)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: Type -> Type -> Type) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (Vector m (Vector n a) -> Matrix m n a
forall (m :: Natural) (n :: Natural) a.
Vector m (Vector n a) -> Matrix m n a
Matrix (Vector m (Vector n a) -> Matrix m n a)
-> ([Vector n a] -> Vector m (Vector n a))
-> [Vector n a]
-> Matrix m n a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Vector n a] -> Vector m (Vector n a)
forall (size :: Natural) a. [a] -> Vector size a
Vector) ([Vector n a], g)
as