-- Copyright (c) David Amos, 2008-2015. All rights reserved.

{-# LANGUAGE MultiParamTypeClasses, TypeSynonymInstances, ScopedTypeVariables, EmptyDataDecls, FlexibleInstances #-}

module Math.Algebra.Field.Extension where

import Prelude hiding ( (<*>) )

import Data.Ratio
import Data.List as L (elemIndex)

import Math.Common.IntegerAsType
import Math.Core.Utils
import Math.Algebra.Field.Base


-- UNIVARIATE POLYNOMIALS

newtype UPoly a = UP [a] deriving (Eq,Ord)
-- the list [a_0, a_1, ..., a_n] represents the polynomial a_0 + a_1 x + ... + a_n x^n

x = UP [0,1] :: UPoly Integer

instance (Eq a, Show a, Num a) => Show (UPoly a) where
    -- show (UP []) = "0"
    show (UP as) = showUP "x" as

showUP _ [] = "0"
showUP v as = let powers = filter ( (/=0) . fst ) $ zip as [0..]
                  c:cs = concatMap showTerm powers
              in if c == '+' then cs else c:cs
    where showTerm (a,i) = showCoeff a ++ showPower a i
          showCoeff a = case show a of
                        "1" -> "+"
                        "-1" -> "-"
                        '-':cs -> '-':cs
                        cs -> '+':cs
          showPower a i | i == 0 = case show a of
                                   "1" -> "1"
                                   "-1" -> "1"
                                   otherwise -> ""
                        | i == 1 = v -- "x"
                        | i > 1  = v ++ "^" ++ show i -- "x^" ++ show i

instance (Eq a, Num a) => Num (UPoly a) where
    UP as + UP bs = UP $ as <+> bs
    negate (UP as) = UP $ map negate as
    UP as * UP bs = UP $ as <*> bs
    fromInteger 0 = UP []
    fromInteger a = UP [fromInteger a]
    abs _ = error "Prelude.Num.abs: inappropriate abstraction"
    signum _ = error "Prelude.Num.signum: inappropriate abstraction"

toUPoly as = UP (reverse (dropWhile (== 0) (reverse as)))

-- The fussiness of the code is to avoid adding trailing zeroes, eg [3] <+> [-3]
-- Otherwise we would have to normalise after every addition
as <+> [] = as
[] <+> bs = bs
-- (a:as) <+> (b:bs) = (a+b) : (as <+> bs)
(a:as) <+> (b:bs) = let c = a+b
                        cs = as <+> bs
                    in if c == 0 && null cs then [] else c:cs

-- The fussiness of the code is to avoid adding trailing zeroes.
-- Note that since we call <+>, we rely on it having similar properties.
[] <*> _ = []
_ <*> [] = [] -- to avoid [0,1] <*> [] -> [0]
(a:as) <*> bs = if null as then map (a*) bs else map (a*) bs <+> (0 : as <*> bs)

-- > let valid xs = null xs || last xs /= 0
-- > quickCheck (\as bs -> not (valid as) || not (valid bs) || valid (as <*> bs))

{-
-- The following definition introduces unnecessary trailing zeroes, eg [3] <*> [2] -> [6,0]
[] <*> _ = []
_ <*> [] = []
(a:as) <*> (b:bs) = [a*b] <+> (0 : map (a*) bs) <+> (0 : map (*b) as) <+> (0 : 0 : as <*> bs)
-}

convert (UP as) = toUPoly $ map fromInteger as
-- Can be used with type annotations to construct polynomials over other types, eg
-- > convert (x^2+3*x+2) :: UPoly F2
-- x^2+x
-- > convert (x^2+3*x+2) :: UPoly F3
-- x^2+2


-- DIVISION ALGORITHM

-- degree
deg (UP as) = length as

-- leading term
lt (UP as) = last as

monomial a i = UP $ replicate i 0 ++ [a]

-- quotRem for UPolys over a field
quotRemUP :: (Eq k, Fractional k) => UPoly k -> UPoly k -> (UPoly k, UPoly k)
quotRemUP f g = qr 0 f where
    qr q r = if deg r < deg_g
             then (q,r)
             else let s = monomial (lt r / lt_g) (deg r - deg_g)
                  in qr (q+s) (r-s*g)
    deg_g = deg g
    lt_g = lt g


modUP f g = snd $ quotRemUP f g

-- extendedEuclidUP f g returns (u,v,d) such that u*f + v*g = d
extendedEuclidUP f g = extendedEuclidUP' f g [] where
    extendedEuclidUP' d 0 qs = let (u,v) = unwind 1 0 qs in (u,v,d)
    extendedEuclidUP' f g qs = let (q,r) = quotRemUP f g in extendedEuclidUP' g r (q:qs)
    unwind u v [] = (u,v)
    unwind u v (q:qs) = unwind v (u-v*q) qs


-- EXTENSION FIELDS

class PolynomialAsType k poly where
    pvalue :: (k,poly) -> UPoly k

data ExtensionField k poly = Ext (UPoly k) deriving (Eq,Ord)

instance (Eq k, Show k, Num k) => Show (ExtensionField k poly) where
    -- show (Ext f) = show f
    -- show (Ext (UP [])) = "0"
    show (Ext (UP as)) = showUP "a" as

instance (Eq k, Fractional k, PolynomialAsType k poly) => Num (ExtensionField k poly) where
    Ext x + Ext y = Ext $ (x+y) -- `modUP` pvalue (undefined :: (k,poly))
    Ext x * Ext y = Ext $ (x*y) `modUP` pvalue (undefined :: (k,poly))
    negate (Ext x) = Ext $ negate x
    fromInteger x = Ext $ fromInteger x
    abs _ = error "Prelude.Num.abs: inappropriate abstraction"
    signum _ = error "Prelude.Num.signum: inappropriate abstraction"

instance (Eq k, Fractional k, PolynomialAsType k poly) => Fractional (ExtensionField k poly) where
    recip 0 = error "ExtensionField.recip 0"
    recip (Ext f) = let g = pvalue (undefined :: (k,poly))
                        (u,v,d@(UP [c])) = extendedEuclidUP f g
                        -- so u*f + v*g == d. We know the d is a unit, ie field element, since g is irreducible
                    in Ext $ (c /> u) `modUP` g
    fromRational q = fromInteger a / fromInteger b where a = numerator q; b = denominator q

-- divide through
c /> f@(UP as) | c == 1 = f
               | c /= 0 = UP (map (c' *) as) where c' = recip c

instance (FiniteField k, PolynomialAsType k poly) => FiniteField (ExtensionField k poly) where
    eltsFq _ = map Ext (polys (d-1) fp) where
        fp = eltsFq (undefined :: k)
        d = deg $ pvalue (undefined :: (k,poly))
    basisFq _ = map embed $ take (d-1) $ iterate (*x) 1 where
        d = deg $ pvalue (undefined :: (k,poly))

-- Not sure if Eq fp is required, need to check with ghc >= 7.4.1
instance (FinSet fp, Eq fp, Num fp, PolynomialAsType fp poly) => FinSet (ExtensionField fp poly) where
    elts = map Ext (polys (d-1) fp') where
        fp' = elts
        d = deg $ pvalue (undefined :: (fp,poly))

embed f = Ext (convert f)


-- PRIME POWER FINITE FIELDS

polys d fp = map toUPoly $ polys' d where
    polys' 0 = [[]]
    polys' d = [x:xs | x <- fp, xs <- polys' (d-1)] -- return in ascending order
    -- polys' d = [x:xs | xs <- polys' (d-1), x <- fp] -- return with elts of fp first

-- Conway polynomials from Holt, Handbook of Computational Group Theory, p60

data ConwayF4
instance PolynomialAsType F2 ConwayF4 where pvalue _ = convert $ x^2+x+1
type F4 = ExtensionField F2 ConwayF4
f4 = map Ext (polys 2 f2) :: [F4]
a4 = embed x :: F4

data ConwayF8
instance PolynomialAsType F2 ConwayF8 where pvalue _ = convert $ x^3+x+1
type F8 = ExtensionField F2 ConwayF8
f8 = map Ext (polys 3 f2) :: [F8]
a8 = embed x :: F8

data ConwayF9
instance PolynomialAsType F3 ConwayF9 where pvalue _ = convert $ x^2+2*x+2
type F9 = ExtensionField F3 ConwayF9
f9 = map Ext (polys 2 f3) :: [F9]
a9 = embed x :: F9

data ConwayF16
instance PolynomialAsType F2 ConwayF16 where pvalue _ = convert $ x^4+x+1
type F16 = ExtensionField F2 ConwayF16
f16 = map Ext (polys 4 f2) :: [F16]
a16 = embed x :: F16

data ConwayF25
instance PolynomialAsType F5 ConwayF25 where pvalue _ = convert $ x^2+4*x+2
type F25 = ExtensionField F5 ConwayF25
f25 = map Ext (polys 2 f5) :: [F25]
a25 = embed x :: F25

data ConwayF27
instance PolynomialAsType F3 ConwayF27 where pvalue _ = convert $ x^3+2*x+1
type F27 = ExtensionField F3 ConwayF27
f27 = map Ext (polys 3 f3) :: [F27]
a27 = embed x :: F27

data ConwayF32
instance PolynomialAsType F2 ConwayF32 where pvalue _ = convert $ x^5+x^2+1
type F32 = ExtensionField F2 ConwayF32
f32 = map Ext (polys 5 f2) :: [F32]
a32 = embed x :: F32


-- generator for the automorphism group of fq, as applied to an element of fq
frobeniusAut x = x ^ p where
    p = char $ eltsFq x

-- the degree of fq as an extension over fp
-- (hence also, the order of the automorphism group of fq)
degree fq = n where
    q = length fq
    p = char fq
    Just n = L.elemIndex q $ iterate (*p) 1



-- QUADRATIC EXTENSIONS OF Q

data Sqrt a = Sqrt a

-- n should be square-free
instance IntegerAsType n => PolynomialAsType Q (Sqrt n) where
    pvalue _ = convert $ x^2 - fromInteger (value (undefined :: n))

type QSqrt2 = ExtensionField Q (Sqrt T2)
sqrt2 = embed x :: QSqrt2

type QSqrt3 = ExtensionField Q (Sqrt T3)
sqrt3 = embed x :: QSqrt3

type QSqrt5 = ExtensionField Q (Sqrt T5)
sqrt5 = embed x :: QSqrt5

type QSqrt7 = ExtensionField Q (Sqrt T7)
sqrt7 = embed x :: QSqrt7

type QSqrtMinus1 = ExtensionField Q (Sqrt TMinus1)
i = embed x :: QSqrtMinus1


type QSqrtMinus2 = ExtensionField Q (Sqrt (M TMinus1 T2))
sqrtminus2 = embed x :: QSqrtMinus2

type QSqrtMinus3 = ExtensionField Q (Sqrt (M TMinus1 T3))
sqrtminus3 = embed x :: QSqrtMinus3

type QSqrtMinus5 = ExtensionField Q (Sqrt (M TMinus1 T5))
sqrtminus5 = embed x :: QSqrtMinus5


-- conjugation automorphism of quadratic field
-- conjugate of a + b sqrt d is a - b sqrt d
conjugate :: ExtensionField Q (Sqrt d) -> ExtensionField Q (Sqrt d)
conjugate (Ext (UP [a,b])) = Ext (UP [a,-b])
conjugate x = x -- the zero or constant cases