module ExtensionField
( ExtensionField
, PolynomialRing
, IrreducibleMonic(split)
, fromField
, toField
, pattern X
, pattern X2
, pattern X3
, pattern Y
) where

import Protolude as P hiding (Semiring, quot, quotRem, rem)

import Data.Euclidean (Euclidean(..), GcdDomain(..))
import Data.Poly.Semiring (VPoly, leading, monomial, scale, toPoly, unPoly, pattern X)
import Data.Semiring as S (Ring(..), Semiring(..))
import Data.Vector (fromList)
import Test.Tasty.QuickCheck (Arbitrary(..), vector)
import Text.PrettyPrint.Leijen.Text (Pretty(..))

import GaloisField (Field(..), GaloisField(..))

-------------------------------------------------------------------------------
-- Data types
-------------------------------------------------------------------------------

-- | Extension fields @GF(p^q)[X]/\<f(X)\>@ for @p@ prime, @q@ positive, and
-- @f(X)@ irreducible monic in @GF(p^q)[X]@.
newtype ExtensionField k im = EF (VPoly k)
deriving (Eq, Generic, Ord, Show)

-- | Polynomial rings.
type PolynomialRing = VPoly

-- | Irreducible monic splitting polynomial @f(X)@ of extension field.
class GaloisField k => IrreducibleMonic k im where
{-# MINIMAL split #-}
-- | Splitting polynomial @f(X)@.
split :: ExtensionField k im -> VPoly k
-- Splitting polynomial degree.
deg' :: ExtensionField k im -> Int
deg' = pred . fromIntegral . degree . split
{-# INLINABLE deg' #-}

-- Extension fields are Galois fields.
instance IrreducibleMonic k im => GaloisField (ExtensionField k im) where
char = const (char (witness :: k))
{-# INLINABLE char #-}
deg  = (deg (witness :: k) *) . deg'
{-# INLINABLE deg #-}
frob = pow <*> char
{-# INLINABLE frob #-}

{-# RULES "ExtensionField/pow"
forall (k :: IrreducibleMonic k im => ExtensionField k im) n . (^) k n = pow k n
#-}

-------------------------------------------------------------------------------
-- Numeric instances
-------------------------------------------------------------------------------

-- Extension fields are fractional.
instance IrreducibleMonic k im => Fractional (ExtensionField k im) where
recip (EF x)        = EF (polyInv x (split (witness :: ExtensionField k im)))
{-# INLINABLE recip #-}
fromRational (x:%y) = fromInteger x / fromInteger y
{-# INLINABLE fromRational #-}

-- Extension fields are numeric.
instance IrreducibleMonic k im => Num (ExtensionField k im) where
EF x + EF y   = EF (plus x y)
{-# INLINE (+) #-}
EF x * EF y   = EF (rem (times x y) (split (witness :: ExtensionField k im)))
{-# INLINABLE (*) #-}
EF x - EF y   = EF (x - y)
{-# INLINE (-) #-}
negate (EF x) = EF (S.negate x)
{-# INLINE negate #-}
fromInteger   = EF . fromInteger
{-# INLINABLE fromInteger #-}
abs           = panic "not implemented."
signum        = panic "not implemented."

-------------------------------------------------------------------------------
-- Semiring instances
-------------------------------------------------------------------------------

-- Extension fields are Euclidean domains.
instance IrreducibleMonic k im => Euclidean (ExtensionField k im) where
quotRem = (flip (,) 0 .) . (/)
{-# INLINE quotRem #-}
degree  = panic "not implemented."

-- Extension fields are fields.
instance IrreducibleMonic k im => Field (ExtensionField k im) where
invert = recip
{-# INLINE invert #-}
minus  = (-)
{-# INLINE minus #-}

-- Extension fields are GCD domains.
instance IrreducibleMonic k im => GcdDomain (ExtensionField k im)

-- Extension fields are rings.
instance IrreducibleMonic k im => Ring (ExtensionField k im) where
negate = P.negate
{-# INLINE negate #-}

-- Extension fields are semirings.
instance IrreducibleMonic k im => Semiring (ExtensionField k im) where
zero        = 0
{-# INLINE zero #-}
plus        = (+)
{-# INLINE plus #-}
one         = 1
{-# INLINE one #-}
times       = (*)
{-# INLINE times #-}
fromNatural = fromIntegral
{-# INLINABLE fromNatural #-}

-------------------------------------------------------------------------------
-- Other instances
-------------------------------------------------------------------------------

-- Extension fields are arbitrary.
instance IrreducibleMonic k im => Arbitrary (ExtensionField k im) where
arbitrary = toField <\$> vector (deg' (witness :: ExtensionField k im))
{-# INLINABLE arbitrary #-}

-- Extension fields are pretty.
instance IrreducibleMonic k im => Pretty (ExtensionField k im) where
pretty (EF x) = pretty (toList (unPoly x))

-- Extension fields are random.
instance IrreducibleMonic k im => Random (ExtensionField k im) where
random  = first toField . unfold (deg' (witness :: ExtensionField k im)) []
where
unfold n xs g
| n <= 0    = (xs, g)
| otherwise = case random g of
(x, g') -> unfold (n - 1) (x : xs) g'
{-# INLINABLE random #-}
randomR = panic "not implemented."

-------------------------------------------------------------------------------
-- Type conversions
-------------------------------------------------------------------------------

-- | Convert from field element to list representation.
fromField :: ExtensionField k im -> [k]
fromField (EF x) = toList (unPoly x)
{-# INLINABLE fromField #-}

-- | Convert from list representation to field element.
toField :: forall k im . IrreducibleMonic k im => [k] -> ExtensionField k im
toField = EF . flip rem (split (witness :: ExtensionField k im)) . toPoly . fromList
{-# INLINABLE toField #-}

-- | Pattern for @X^2@.
pattern X2 :: GaloisField k => VPoly k
pattern X2 <- _ where X2 = toPoly (fromList [0, 0, 1])

-- | Pattern for @X^3@.
pattern X3 :: GaloisField k => VPoly k
pattern X3 <- _ where X3 = toPoly (fromList [0, 0, 0, 1])

-- | Pattern for descending tower of indeterminate variables.
pattern Y :: IrreducibleMonic k im => VPoly k -> VPoly (ExtensionField k im)
pattern Y <- _ where Y = monomial 0 . EF

-------------------------------------------------------------------------------
-- Polynomial arithmetic
-------------------------------------------------------------------------------

-- Polynomial inversion algorithm.
polyInv :: GaloisField k => VPoly k -> VPoly k -> VPoly k
polyInv xs ps = case first leading (polyGCD xs ps) of
(Just (0, x), ys) -> scale 0 (recip x) ys
_                 -> panic "no multiplicative inverse."
{-# INLINABLE polyInv #-}

-- Polynomial extended greatest common divisor algorithm.
polyGCD :: forall k . GaloisField k => VPoly k -> VPoly k -> (VPoly k, VPoly k)
polyGCD x y = polyGCD' 0 1 y x
where
polyGCD' :: VPoly k -> VPoly k -> VPoly k -> VPoly k -> (VPoly k, VPoly k)
polyGCD' s _  r 0  = (r, s)
polyGCD' s s' r r' = case quot r r' of
q -> polyGCD' s' (s - times q s') r' (r - times q r')
{-# INLINABLE polyGCD #-}