-- | Type class interface to different implementations of finite fields

{-# LANGUAGE BangPatterns, FlexibleContexts, TypeApplications #-}
{-# LANGUAGE ScopedTypeVariables, TypeFamilies, TypeFamilyDependencies #-}
{-# LANGUAGE ExistentialQuantification, StandaloneDeriving #-}
{-# LANGUAGE DataKinds #-}

module Math.FiniteField.Class where

--------------------------------------------------------------------------------

import Data.Bits
import Data.List

import GHC.TypeNats (Nat)

import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map

import System.Random ( RandomGen )

import Math.FiniteField.TypeLevel
import Math.FiniteField.TypeLevel.Singleton

--------------------------------------------------------------------------------
-- * Fields

{-
-- NOTE: recursive classes just cause problems

-- | A class for witness types (these witness the existence of a field)
class 
  ( Show w, Field (FieldElem w)
  , WitnessPrime w ~ Prime (FieldElem w) 
  , WitnessDim   w ~ Dim   (FieldElem w) 
  ) => FieldWitness w 
  where
    type FieldElem    w = f | f -> w
    type WitnessPrime w :: Nat
    type WitnessDim   w :: Nat
-}

-- | A class for field element types 
class (Eq f, Ord f, Show f, Num f, Fractional f, Show (Witness f)) => Field f where
    -- | witness for the existence of the field (this is an injective type family!) 
    type Witness f = w | w -> f   
    -- | the characteristic at type level
    type Prime f :: Nat
    -- | the dimension at type level
    type Dim f :: Nat              
    -- | the prime characteristic
    characteristic   :: Witness f -> Integer   
    -- | dimension over the prime field (the exponent @m@ in @q=p^m@)
    dimension        :: Witness f -> Integer    -- TODO: this should be Int
    -- | the size (or order) of the field
    fieldSize        :: Witness f -> Integer   
    -- | The additive identity of the field
    zero             :: Witness f -> f
    -- | The multiplicative identity of the field
    one              :: Witness f -> f
    -- | check for equality with the additive identity
    isZero           :: f -> Bool
    -- | check for equality with the multiplicative identity
    isOne            :: f -> Bool
    -- | an element of the prime field
    embed            :: Witness f -> Integer -> f         
    embedSmall       :: Witness f -> Int     -> f   
    -- | a uniformly random field element
    randomFieldElem  :: RandomGen gen => Witness f -> gen -> (f,gen) 
    -- | a random invertible element
    randomInvertible :: RandomGen gen => Witness f -> gen -> (f,gen) 
    -- | a primitive generator
    primGen          :: Witness f -> f                    
    -- | extract the witness from a field element
    witnessOf        :: f -> Witness f                    
    -- | exponentiation 
    power            :: f -> Integer -> f                 
    powerSmall       :: f -> Int     -> f            
    -- | Frobenius automorphism @x -> x^p@
    frobenius        :: f -> f
    -- | list of field elements (of course it's only useful for very small fields)
    enumerate        :: Witness f -> [f]                  
  
    -- default implementations
  
    embedSmall !Witness f
w !Int
x = forall f. Field f => Witness f -> Integer -> f
embed Witness f
w (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
x)
    powerSmall !f
x !Int
e = forall f. Field f => f -> Integer -> f
power f
x (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
e)      -- it's important not to use powerDefault here, because 'power' may be more efficient
    fieldSize  !Witness f
w    = forall f. Field f => Witness f -> Integer
characteristic Witness f
w forall a b. (Num a, Integral b) => a -> b -> a
^ forall f. Field f => Witness f -> Integer
dimension Witness f
w
    power            = forall f. Field f => f -> Integer -> f
powerDefault 
    frobenius     !f
x = forall f. Field f => f -> Integer -> f
power f
x (forall f. Field f => Witness f -> Integer
characteristic (forall f. Field f => f -> Witness f
witnessOf f
x))
  
    zero       !Witness f
w    = forall f. Field f => Witness f -> Int -> f
embedSmall Witness f
w Int
0
    one        !Witness f
w    = forall f. Field f => Witness f -> Int -> f
embedSmall Witness f
w Int
1
    -- isZero     !x    = (x == zero w)   -- we don't have a witness available here...
    -- isOne      !x    = (x == one  w)
  
    randomInvertible !Witness f
w !gen
g = case forall f gen.
(Field f, RandomGen gen) =>
Witness f -> gen -> (f, gen)
randomFieldElem Witness f
w gen
g of 
      (f
x,gen
g') -> if forall f. Field f => f -> Bool
isZero f
x then forall f gen.
(Field f, RandomGen gen) =>
Witness f -> gen -> (f, gen)
randomInvertible Witness f
w gen
g' else (f
x,gen
g')

fieldPrimeSNat :: Field f => Witness f -> SNat (Prime f)
fieldPrimeSNat :: forall f. Field f => Witness f -> SNat (Prime f)
fieldPrimeSNat Witness f
w = forall (n :: Nat). Integer -> SNat n
SNat (forall f. Field f => Witness f -> Integer
characteristic Witness f
w)

fieldPrimeSNat64 :: Field f => Witness f -> SNat64 (Prime f)
fieldPrimeSNat64 :: forall f. Field f => Witness f -> SNat64 (Prime f)
fieldPrimeSNat64 Witness f
w = forall (n :: Nat). Word64 -> SNat64 n
SNat64 (forall a. Num a => Integer -> a
fromInteger forall a b. (a -> b) -> a -> b
$ forall f. Field f => Witness f -> Integer
characteristic Witness f
w)

fieldDimSNat :: Field f => Witness f -> SNat (Dim f)
fieldDimSNat :: forall f. Field f => Witness f -> SNat (Dim f)
fieldDimSNat Witness f
w = forall (n :: Nat). Integer -> SNat n
SNat (forall f. Field f => Witness f -> Integer
dimension Witness f
w)

fieldDimSNat64 :: Field f => Witness f -> SNat64 (Dim f)
fieldDimSNat64 :: forall f. Field f => Witness f -> SNat64 (Dim f)
fieldDimSNat64 Witness f
w = forall (n :: Nat). Word64 -> SNat64 n
SNat64 (forall a. Num a => Integer -> a
fromInteger forall a b. (a -> b) -> a -> b
$ forall f. Field f => Witness f -> Integer
dimension Witness f
w)

--------------------------------------------------------------------------------

data SomeField 
  = forall f. Field f => SomeField (Witness f)

deriving instance Show SomeField

--------------------------------------------------------------------------------

-- -- * Subfields
-- 
-- class 
--   ( FieldWitness (AmbientWitness w)
--   , FieldWitness (SubWitness     w) 
--   ) => SubFieldWitness w 
--   where
--     type AmbientWitness w :: *
--     type SubWitness     w :: *
-- 
-- class (Field (AmbientField s) , Field (SubField s)) => SubField s where
--   type AmbientField s :: *
--   type SubField     s :: *
--   type SubFieldWitness 
--   embedSubField :: SubField s -> AmbientField s 
--   mbSubField    :: AmbientField s -> Maybe (SubField s)
--   isSubField    :: AmbientField s -> Bool
--   coordinates   :: AmbientField s -> [SubField s]

--------------------------------------------------------------------------------
-- * Some generic functions

-- | Returns @"GF(p)"@ or @"GF(p^m)"@
fieldName :: Field f => Witness f -> String
fieldName :: forall f. Field f => Witness f -> String
fieldName Witness f
w 
  | Integer
m forall a. Eq a => a -> a -> Bool
== Integer
1     = String
"GF(" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Integer
p forall a. [a] -> [a] -> [a]
++ String
")"
  | Bool
otherwise  = String
"GF(" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Integer
p forall a. [a] -> [a] -> [a]
++ String
"^" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Integer
m forall a. [a] -> [a] -> [a]
++ String
")"
  where
    p :: Integer
p = forall f. Field f => Witness f -> Integer
characteristic Witness f
w
    m :: Integer
m = forall f. Field f => Witness f -> Integer
dimension      Witness f
w

-- | The multiplicate inverse (synonym for 'recip')
inverse :: Field f => f -> f
inverse :: forall f. Field f => f -> f
inverse = forall a. Fractional a => a -> a
recip

-- | Enumerate the elements of the prime field only 
enumPrimeField :: forall f. Field f => Witness f -> [f]
enumPrimeField :: forall f. Field f => Witness f -> [f]
enumPrimeField Witness f
w = [ forall f. Field f => Witness f -> Int -> f
embedSmall Witness f
w Int
i | Int
i<-[Int
0..Int
pforall a. Num a => a -> a -> a
-Int
1] ] where
  pbig :: Integer
pbig = forall f. Field f => Witness f -> Integer
characteristic Witness f
w
  p :: Int
p    = forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
pbig :: Int

-- | The nonzero elements in cyclic order, starting from the primitive generator
-- (of course this is only useful for very small fields)
multGroup :: Field f => Witness f -> [f]    
multGroup :: forall f. Field f => Witness f -> [f]
multGroup Witness f
w = forall a. (a -> a -> a) -> [a] -> [a]
scanl1 forall a. Num a => a -> a -> a
(*) [f]
list where
  g :: f
g    = forall f. Field f => Witness f -> f
primGen Witness f
w
  m :: Integer
m    = forall f. Field f => Witness f -> Integer
fieldSize Witness f
w
  list :: [f]
list = forall a. Int -> a -> [a]
replicate (forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
m forall a. Num a => a -> a -> a
- Int
1) f
g

--------------------------------------------------------------------------------

-- | Computes a table of discrete logarithms with respect to the primitive 
-- generator. Note: zero (that is, the additive identitiy of the field) 
-- is not present in the resulting map.
discreteLogTable :: forall f. Field f => Witness f -> Map f Int
discreteLogTable :: forall f. Field f => Witness f -> Map f Int
discreteLogTable Witness f
witness = forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList (Int -> f -> [(f, Int)]
worker Int
0 (forall f. Field f => Witness f -> f
one Witness f
witness)) where
  g :: f
g = forall f. Field f => Witness f -> f
primGen   Witness f
witness
  q :: Integer
q = forall f. Field f => Witness f -> Integer
fieldSize Witness f
witness
  qm1 :: Int
qm1 = forall a. Num a => Integer -> a
fromInteger Integer
q forall a. Num a => a -> a -> a
- Int
1
  worker :: Int -> f -> [(f,Int)]
  worker :: Int -> f -> [(f, Int)]
worker !Int
e !f
acc
    | Int
e forall a. Ord a => a -> a -> Bool
< Int
qm1    = (f
acc,Int
e) forall a. a -> [a] -> [a]
: Int -> f -> [(f, Int)]
worker (Int
eforall a. Num a => a -> a -> a
+Int
1) (f
accforall a. Num a => a -> a -> a
*f
g)
    | Bool
otherwise  = []

--------------------------------------------------------------------------------

-- | Generic exponentiation
powerDefault :: forall f. (Field f) => f -> Integer -> f
powerDefault :: forall f. Field f => f -> Integer -> f
powerDefault !f
z !Integer
e 
  | forall f. Field f => f -> Bool
isZero f
z  = f
z 
  | Integer
e forall a. Eq a => a -> a -> Bool
== Integer
0    = forall f. Field f => Witness f -> f
one Witness f
w
  | Integer
e forall a. Ord a => a -> a -> Bool
< Integer
0     = forall f. Field f => f -> Integer -> f
powerDefault (forall a. Fractional a => a -> a
recip f
z) (forall a. Num a => a -> a
negate Integer
e)
  | Integer
e forall a. Ord a => a -> a -> Bool
>= Integer
pm1  = f -> f -> Integer -> f
go (forall f. Field f => Witness f -> f
one Witness f
w) f
z (forall a. Integral a => a -> a -> a
mod Integer
e Integer
pm1)
  | Bool
otherwise = f -> f -> Integer -> f
go (forall f. Field f => Witness f -> f
one Witness f
w) f
z Integer
e
  where
    w :: Witness f
w   = forall f. Field f => f -> Witness f
witnessOf f
z
    pm1 :: Integer
pm1 = forall f. Field f => Witness f -> Integer
fieldSize Witness f
w forall a. Num a => a -> a -> a
- Integer
1
    go :: f -> f -> Integer -> f
    go :: f -> f -> Integer -> f
go !f
acc !f
y !Integer
e = if Integer
e forall a. Eq a => a -> a -> Bool
== Integer
0 
      then f
acc
      else case (Integer
e forall a. Bits a => a -> a -> a
.&. Integer
1) of
        Integer
0 -> f -> f -> Integer -> f
go  f
acc    (f
yforall a. Num a => a -> a -> a
*f
y) (forall a. Bits a => a -> Int -> a
shiftR Integer
e Int
1)
        Integer
_ -> f -> f -> Integer -> f
go (f
accforall a. Num a => a -> a -> a
*f
y) (f
yforall a. Num a => a -> a -> a
*f
y) (forall a. Bits a => a -> Int -> a
shiftR Integer
e Int
1)

--------------------------------------------------------------------------------