{-# LANGUAGE OverloadedLists  #-}
{-# LANGUAGE TypeApplications #-}

module ZkFold.Base.Algebra.Basic.Permutations (
    IndexSet,
    IndexPartition,
    Permutation,
    fromPermutation,
    applyPermutation,
    mkIndexPartition,
    fromCycles
) where

import           Data.Functor.Rep                 (Representable (index))
import           Data.Map                         (Map, elems, empty, singleton, union)
import           Data.Maybe                       (fromJust)
import qualified Data.Vector                      as V
import           Prelude                          hiding (Num (..), drop, length, mod, (!!))
import qualified Prelude                          as P
import           Test.QuickCheck                  (Arbitrary (..))

import           ZkFold.Base.Algebra.Basic.Class
import           ZkFold.Base.Algebra.Basic.Number
import           ZkFold.Base.Data.Vector          (Vector (..), fromVector, toVector, unsafeToVector)
import           ZkFold.Prelude                   (chooseNatural, drop, length, (!!))

-- TODO (Issue #18): make the code safer

------------------------------ Index sets and partitions -------------------------------------

type IndexSet = V.Vector Natural
type IndexPartition a = Map a IndexSet

mkIndexPartition :: Ord a => V.Vector a -> IndexPartition a
mkIndexPartition :: forall a. Ord a => Vector a -> IndexPartition a
mkIndexPartition Vector a
vs =
    let f :: a -> Map a (Vector Natural)
f a
i = a -> Vector Natural -> Map a (Vector Natural)
forall k a. k -> a -> Map k a
singleton a
i (Vector Natural -> Map a (Vector Natural))
-> Vector Natural -> Map a (Vector Natural)
forall a b. (a -> b) -> a -> b
$ ((a, Natural) -> Natural) -> Vector (a, Natural) -> Vector Natural
forall a b. (a -> b) -> Vector a -> Vector b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap (a, Natural) -> Natural
forall a b. (a, b) -> b
snd (Vector (a, Natural) -> Vector Natural)
-> Vector (a, Natural) -> Vector Natural
forall a b. (a -> b) -> a -> b
$ ((a, Natural) -> Bool)
-> Vector (a, Natural) -> Vector (a, Natural)
forall a. (a -> Bool) -> Vector a -> Vector a
V.filter (\(a
v, Natural
_) -> a
v a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
i) (Vector (a, Natural) -> Vector (a, Natural))
-> Vector (a, Natural) -> Vector (a, Natural)
forall a b. (a -> b) -> a -> b
$ Vector a -> Vector Natural -> Vector (a, Natural)
forall a b. Vector a -> Vector b -> Vector (a, b)
V.zip Vector a
vs [Natural
Item (Vector Natural)
1 .. Vector a -> Natural
forall (t :: Type -> Type) a. Foldable t => t a -> Natural
length Vector a
vs]
    in (Map a (Vector Natural)
 -> Map a (Vector Natural) -> Map a (Vector Natural))
-> Map a (Vector Natural)
-> Vector (Map a (Vector Natural))
-> Map a (Vector Natural)
forall a b. (a -> b -> a) -> a -> Vector b -> a
V.foldl Map a (Vector Natural)
-> Map a (Vector Natural) -> Map a (Vector Natural)
forall k a. Ord k => Map k a -> Map k a -> Map k a
union Map a (Vector Natural)
forall k a. Map k a
empty (Vector (Map a (Vector Natural)) -> Map a (Vector Natural))
-> Vector (Map a (Vector Natural)) -> Map a (Vector Natural)
forall a b. (a -> b) -> a -> b
$ (a -> Map a (Vector Natural))
-> Vector a -> Vector (Map a (Vector Natural))
forall a b. (a -> b) -> Vector a -> Vector b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> Map a (Vector Natural)
f Vector a
vs

------------------------------------- Permutations -------------------------------------------

newtype Permutation n = Permutation (Vector n Natural)
    deriving (Int -> Permutation n -> ShowS
[Permutation n] -> ShowS
Permutation n -> String
(Int -> Permutation n -> ShowS)
-> (Permutation n -> String)
-> ([Permutation n] -> ShowS)
-> Show (Permutation n)
forall (n :: Natural). Int -> Permutation n -> ShowS
forall (n :: Natural). [Permutation n] -> ShowS
forall (n :: Natural). Permutation n -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall (n :: Natural). Int -> Permutation n -> ShowS
showsPrec :: Int -> Permutation n -> ShowS
$cshow :: forall (n :: Natural). Permutation n -> String
show :: Permutation n -> String
$cshowList :: forall (n :: Natural). [Permutation n] -> ShowS
showList :: [Permutation n] -> ShowS
Show, Permutation n -> Permutation n -> Bool
(Permutation n -> Permutation n -> Bool)
-> (Permutation n -> Permutation n -> Bool) -> Eq (Permutation n)
forall (n :: Natural). Permutation n -> Permutation n -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall (n :: Natural). Permutation n -> Permutation n -> Bool
== :: Permutation n -> Permutation n -> Bool
$c/= :: forall (n :: Natural). Permutation n -> Permutation n -> Bool
/= :: Permutation n -> Permutation n -> Bool
Eq)

instance KnownNat n => Arbitrary (Permutation n) where
    arbitrary :: Gen (Permutation n)
arbitrary =
        let f :: [a] -> [a] -> Gen [a]
f [a]
as [] = [a] -> Gen [a]
forall a. a -> Gen a
forall (m :: Type -> Type) a. Monad m => a -> m a
return [a]
as
            f [a]
as [a]
bs = do
                Natural
i <- (Natural, Natural) -> Gen Natural
chooseNatural (Natural
0, [a] -> Natural
forall (t :: Type -> Type) a. Foldable t => t a -> Natural
length [a]
bs Natural -> Natural -> Natural
-! Natural
1)
                let as' :: [a]
as' = ([a]
bs [a] -> Natural -> a
forall a. [a] -> Natural -> a
!! Natural
i) a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
as
                    bs' :: [a]
bs' = Natural -> [a] -> [a]
forall a. Natural -> [a] -> [a]
drop Natural
i [a]
bs
                [a] -> [a] -> Gen [a]
f [a]
as' [a]
bs'
        in Vector n Natural -> Permutation n
forall (n :: Natural). Vector n Natural -> Permutation n
Permutation (Vector n Natural -> Permutation n)
-> ([Natural] -> Vector n Natural) -> [Natural] -> Permutation n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Natural] -> Vector n Natural
forall (size :: Natural) a. [a] -> Vector size a
unsafeToVector ([Natural] -> Permutation n)
-> Gen [Natural] -> Gen (Permutation n)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$>
              [Natural] -> [Natural] -> Gen [Natural]
forall {a}. [a] -> [a] -> Gen [a]
f [] [Natural -> Natural
forall a b. FromConstant a b => a -> b
fromConstant Natural
x | Natural
x <- [Natural
Item [Natural]
1..forall (n :: Natural). KnownNat n => Natural
value @n]]

fromPermutation :: Permutation n -> [Natural]
fromPermutation :: forall (n :: Natural). Permutation n -> [Natural]
fromPermutation (Permutation Vector n Natural
perm) = Vector n Natural -> [Natural]
forall (size :: Natural) a. Vector size a -> [a]
fromVector Vector n Natural
perm

applyPermutation :: KnownNat n => Permutation n -> Vector n a -> Vector n a
applyPermutation :: forall (n :: Natural) a.
KnownNat n =>
Permutation n -> Vector n a -> Vector n a
applyPermutation (Permutation Vector n Natural
ps) Vector n a
as =
    let
        -- 1-indexed Natural to 0-indexed Zp n
        naturalToZp :: a -> a
naturalToZp a
ix = a -> a
forall a b. FromConstant a b => a -> b
fromConstant a
ix a -> a -> a
forall a. AdditiveGroup a => a -> a -> a
- a
1
    in
        (Natural -> a) -> Vector n Natural -> Vector n a
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 (Vector n a -> Rep (Vector n) -> a
forall a. Vector n a -> Rep (Vector n) -> a
forall (f :: Type -> Type) a. Representable f => f a -> Rep f -> a
index Vector n a
as (Zp n -> a) -> (Natural -> Zp n) -> Natural -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Natural -> Zp n
forall {a} {a}.
(AdditiveGroup a, Num a, FromConstant a a) =>
a -> a
naturalToZp) Vector n Natural
ps

applyCycle :: V.Vector Natural -> Permutation n -> Permutation n
applyCycle :: forall (n :: Natural).
Vector Natural -> Permutation n -> Permutation n
applyCycle Vector Natural
c (Permutation Vector n Natural
perm) = Vector n Natural -> Permutation n
forall (n :: Natural). Vector n Natural -> Permutation n
Permutation (Vector n Natural -> Permutation n)
-> Vector n Natural -> Permutation n
forall a b. (a -> b) -> a -> b
$ (Natural -> Natural) -> Vector n Natural -> Vector n Natural
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 Natural -> Natural
f Vector n Natural
perm
    where
        f :: Natural -> Natural
        f :: Natural -> Natural
f Natural
i = case Natural
i Natural -> Vector Natural -> Maybe Int
forall a. Eq a => a -> Vector a -> Maybe Int
`V.elemIndex` Vector Natural
c of
            Just Int
j  -> Vector Natural
c Vector Natural -> Int -> Natural
forall a. Vector a -> Int -> a
V.! ((Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
P.+ Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`P.mod` Vector Natural -> Int
forall a. Vector a -> Int
V.length Vector Natural
c)
            Maybe Int
Nothing -> Natural
i

fromCycles :: KnownNat n => IndexPartition a -> Permutation n
fromCycles :: forall (n :: Natural) a.
KnownNat n =>
IndexPartition a -> Permutation n
fromCycles IndexPartition a
p =
    let n :: Natural
n = Int -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Natural) -> Int -> Natural
forall a b. (a -> b) -> a -> b
$ Vector Natural -> Int
forall a. Vector a -> Int
V.length (Vector Natural -> Int) -> Vector Natural -> Int
forall a b. (a -> b) -> a -> b
$ [Vector Natural] -> Vector Natural
forall a. [Vector a] -> Vector a
V.concat ([Vector Natural] -> Vector Natural)
-> [Vector Natural] -> Vector Natural
forall a b. (a -> b) -> a -> b
$ IndexPartition a -> [Vector Natural]
forall k a. Map k a -> [a]
elems IndexPartition a
p
    in (Vector Natural -> Permutation n -> Permutation n)
-> Permutation n -> [Vector Natural] -> Permutation n
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 Vector Natural -> Permutation n -> Permutation n
forall (n :: Natural).
Vector Natural -> Permutation n -> Permutation n
applyCycle (Vector n Natural -> Permutation n
forall (n :: Natural). Vector n Natural -> Permutation n
Permutation (Vector n Natural -> Permutation n)
-> Vector n Natural -> Permutation n
forall a b. (a -> b) -> a -> b
$ Maybe (Vector n Natural) -> Vector n Natural
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe (Vector n Natural) -> Vector n Natural)
-> Maybe (Vector n Natural) -> Vector n Natural
forall a b. (a -> b) -> a -> b
$ [Natural] -> Maybe (Vector n Natural)
forall (size :: Natural) a.
KnownNat size =>
[a] -> Maybe (Vector size a)
toVector [Natural
Item [Natural]
1 .. Natural
Item [Natural]
n]) ([Vector Natural] -> Permutation n)
-> [Vector Natural] -> Permutation n
forall a b. (a -> b) -> a -> b
$ IndexPartition a -> [Vector Natural]
forall k a. Map k a -> [a]
elems IndexPartition a
p