{-# LANGUAGE UndecidableInstances #-}

module ZkFold.Base.Protocol.Plonkup.Utils where

import           Data.Bifunctor                   (first)
import           Data.Bool                        (bool)
import           Data.Map                         (fromList, insertWith, toList)
import           Prelude                          hiding (Num (..), drop, length, replicate, sum, take, (!!), (/), (^))
import           System.Random                    (RandomGen, mkStdGen, uniformR)

import           ZkFold.Base.Algebra.Basic.Class
import           ZkFold.Base.Algebra.Basic.Number
import           ZkFold.Prelude                   (log2ceiling, replicate)

getParams :: forall a . (Eq a, FiniteField a) => Natural -> (a, a, a)
getParams :: forall a. (Eq a, FiniteField a) => Natural -> (a, a, a)
getParams Natural
n = StdGen -> (a, a, a)
forall g. RandomGen g => g -> (a, a, a)
findK' (StdGen -> (a, a, a)) -> StdGen -> (a, a, a)
forall a b. (a -> b) -> a -> b
$ Int -> StdGen
mkStdGen Int
0
    where
        omega :: a
omega = case forall a. Field a => Natural -> Maybe a
rootOfUnity @a (Natural -> Natural
forall a b. (Integral a, Integral b) => a -> b
log2ceiling Natural
n) of
                  Just a
o -> a
o
                  Maybe a
_      -> [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"impossible"
        hGroup :: [a]
hGroup = (Natural -> a) -> [Natural] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (a
omegaa -> Natural -> a
forall a b. Exponent a b => a -> b -> a
^) [Natural
0 .. Natural
nNatural -> Natural -> Natural
-!Natural
1]
        hGroup' :: a -> [a]
hGroup' a
k = (a -> a) -> [a] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (a
ka -> a -> a
forall a. MultiplicativeSemigroup a => a -> a -> a
*) [a]
hGroup

        findK' :: RandomGen g => g -> (a, a, a)
        findK' :: forall g. RandomGen g => g -> (a, a, a)
findK' g
g =
            let (a
k1, g
g') = (Natural -> a) -> (Natural, g) -> (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 Natural -> a
forall a b. FromConstant a b => a -> b
fromConstant ((Natural, g) -> (a, g)) -> (Natural, g) -> (a, g)
forall a b. (a -> b) -> a -> b
$ (Natural, Natural) -> g -> (Natural, g)
forall g a. (RandomGen g, UniformRange a) => (a, a) -> g -> (a, g)
uniformR (Natural
1, forall a. Finite a => Natural
order @a Natural -> Natural -> Natural
-! Natural
1) g
g
                (a
k2, g
g'') = (Natural -> a) -> (Natural, g) -> (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 Natural -> a
forall a b. FromConstant a b => a -> b
fromConstant ((Natural, g) -> (a, g)) -> (Natural, g) -> (a, g)
forall a b. (a -> b) -> a -> b
$ (Natural, Natural) -> g -> (Natural, g)
forall g a. (RandomGen g, UniformRange a) => (a, a) -> g -> (a, g)
uniformR (Natural
1, forall a. Finite a => Natural
order @a Natural -> Natural -> Natural
-! Natural
1) g
g'
            in (a, a, a) -> (a, a, a) -> Bool -> (a, a, a)
forall a. a -> a -> Bool -> a
bool (g -> (a, a, a)
forall g. RandomGen g => g -> (a, a, a)
findK' g
g'') (a
omega, a
k1, a
k2) (Bool -> (a, a, a)) -> Bool -> (a, a, a)
forall a b. (a -> b) -> a -> b
$
                (a -> Bool) -> [a] -> Bool
forall (t :: Type -> Type) a.
Foldable t =>
(a -> Bool) -> t a -> Bool
all (a -> [a] -> Bool
forall (t :: Type -> Type) a.
(Foldable t, Eq a) =>
a -> t a -> Bool
`notElem` [a]
hGroup) (a -> [a]
hGroup' a
k1)
                Bool -> Bool -> Bool
&& (a -> Bool) -> [a] -> Bool
forall (t :: Type -> Type) a.
Foldable t =>
(a -> Bool) -> t a -> Bool
all (a -> [a] -> Bool
forall (t :: Type -> Type) a.
(Foldable t, Eq a) =>
a -> t a -> Bool
`notElem` a -> [a]
hGroup' a
k1) (a -> [a]
hGroup' a
k2)

sortByList :: Ord a => [a] -> [a] -> [a]
sortByList :: forall a. Ord a => [a] -> [a] -> [a]
sortByList [a]
f [a]
t =
    let m :: Map a Natural
m  = [(a, Natural)] -> Map a Natural
forall k a. Ord k => [(k, a)] -> Map k a
fromList ([(a, Natural)] -> Map a Natural)
-> [(a, Natural)] -> Map a Natural
forall a b. (a -> b) -> a -> b
$ [a] -> [Natural] -> [(a, Natural)]
forall a b. [a] -> [b] -> [(a, b)]
zip [a]
t (forall a. a -> [a]
repeat @Natural Natural
0)
        m' :: Map a Natural
m' = (Map a Natural -> a -> Map a Natural)
-> Map a Natural -> [a] -> Map a Natural
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 (\Map a Natural
acc a
x -> (Natural -> Natural -> Natural)
-> a -> Natural -> Map a Natural -> Map a Natural
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
insertWith Natural -> Natural -> Natural
forall a. AdditiveSemigroup a => a -> a -> a
(+) a
x Natural
1 Map a Natural
acc) Map a Natural
m [a]
f
    in ((a, Natural) -> [a]) -> [(a, Natural)] -> [a]
forall (t :: Type -> Type) a b.
Foldable t =>
(a -> [b]) -> t a -> [b]
concatMap (\(a
k, Natural
v) -> Natural -> a -> [a]
forall a. Natural -> a -> [a]
replicate Natural
v a
k) ([(a, Natural)] -> [a]) -> [(a, Natural)] -> [a]
forall a b. (a -> b) -> a -> b
$ Map a Natural -> [(a, Natural)]
forall k a. Map k a -> [(k, a)]
toList Map a Natural
m'