module BerlekampAlgorithm ( g,frob, pivotPos',lswap,triangulizedModIntegerMat,nullSpaceModIntegerMat,mmultZ,
matrixBerlTranspose ,derivPolyZ,squareFreePolyZ,irreducibilityTestPolyZ,berlekamp,multPoly) where
import Data.List
import Bezout 
-- |Berlekamp's Factorization Algorithm over Fp[x] : computes the factorization of a monic square-free polynomial P into irreducible factor polynomials over F_{p}[x] , p is a prime number. This method is based on linear algebra over finite field.
-- | g
g x = if x == Nothing then (- 1) else (\(Just i)->i) x

-- | pivotPos'
pivotPos' z x = let y = pivotMin x in 
		case y of (-1,-1) -> (-1,-1)
			  _  -> pivotPos z x where 
				pivotPos _ [[]]  = (-1,-1)
				pivotPos z x = let y = map (head) x in 
					let w = (findIndex (/= 0) y) in 
					let w' = g $! w in 
					if (w' /= (-1)) then (w',z) else pivotPos (z+1) (map (tail) $ x) where 
pivotMin x = let y = map (findIndex (/=0)) x in 
		let w = (findIndex (/= Nothing) y) in 
		let z = (find (/= Nothing) y) in (g w,h z) where 
		h x = if x == Nothing then (- 1) else (\(Just (Just i))->i) x

-- | lswap
lswap i j (xs , ys) = (sswap i j xs , sswap i j ys) where 
sswap i j xs = if i == j then xs else 
		take l xs ++ [xs !! u] ++ (take (u-l-1) $ drop (l+1) xs) ++ [xs !! l] ++ drop (u+1) xs
          where l = if i<j then i else j
                u = if i>j then i else j

-- | triangulizedModIntegerMat
--triangulizedModIntegerMat p m: gives the gauss triangular decomposition of an integeral matrix m in Fp.
-- The result is (r, u) where u is a unimodular matrix, r is an upper-triangular matrix , and u.m = r.

triangulizedModIntegerMat p m = let b = fromIntegral $ length m in hnf' p (m , matI b) where
hnf' :: Integer ->  ([[Integer]],[[Integer]]) -> ([[Integer]],[[Integer]])
hnf' p (a,b) =  let (px , py) = pivotPos' 0 a in
		case (px , py) of (-1,-1) -> (a,b)
			          _ -> let (u ,v) = lswap 0 px (a,b) in
				       let (c:cs , d:ds) = gaussZ' p (u,v) in 
				       let (nu , nv) = hnf' p ( cs , ds) in (c:nu , d:nv) where 
gaussZ' p (u,v) = let (x,y) = gaussZ p (u,v) in (last x : init x , last y : init y) where 
gaussZ :: Integer -> ([[Integer]],[[Integer]]) -> ([[Integer]],[[Integer]])
gaussZ _ ([x], a ) = ([x], a)
gaussZ p ((x:y:xs), (s:t:ix)) = let beta = fromIntegral $ g (findIndex (/= 0) x) in 
		let  [a,b] = [x!!beta, y!!beta] in 
		let d = (b * (inverseMod a p)) in
		let r = red p d x y in let m = red p d s t in 
		let (u , v) = gaussZ p ((x:xs) ,(s:ix)) in (r:u, m:v) where 
		red p d x y = zipWith (\c1 c2 -> mods (c2 - (c1 * d)) p ) x y where 
matI n = [ [fromIntegral $ fromEnum $ i == j | i <- [1..n]] | j <- [1..n]]

-- | nullSpaceModIntegerMat p m : computes the null space of matrix m in Fp
nullSpaceModIntegerMat :: Integer -> [[Integer]] -> [[Integer]]
nullSpaceModIntegerMat p x = let (y,z) = triangulizedModIntegerMat p x in ff (y,z) where
ff ([],_) = []
ff (x:xs,y:ys) =  if (nub $! x) == [0] then y : ff (xs,ys) else ff (xs,ys)


-- | mmultZ
-- mmultZ p a b : compute the product of two integer matrices in Fp.
mmultZ p a b = [ [ let y = sum $ zipWith (*) ar bc in mods y p | bc <- (transpose b)] | ar <- a ]


-- | gcdPolyZ
-- gcdPolyZ p P1 P2 : gives the polynomial gcd of P1 , P2 modulo over Fp[x].
gcdPolyZ p x y = let u = last $ extendedgcdpoly p x y in let v = head u in 
			let t = inverseMod v p in mMod (map (* t) u) p

--frob
-- | Frobenius automorphism : linear map V -> V^p - V , V in Fp[x]/P and Fp[x]/P as vector space over the field Fp.

frob :: Integer -> Integer -> [Integer] -> [Integer]
frob _ 0 _ = [0]
frob p k h = let y = prettyFormPoly [[1, p * k],[- 1 , k ],[0,0]] in reverse $! last $! euclideanPolyMod p y h 


-- |matrixBerl
-- matrixBerl p f : is the matrix of the Frobenius endomorphism over the canonical base {1,X,X^2..,X^(p-1)} ,
-- matrixBerl(i,j) = X^(pj)-X^j mod P.
matrixBerlTranspose :: Integer -> [Integer] -> [[Integer]]
matrixBerlTranspose p h = let a = genericLength h - 1 in
	 [let u = fromIntegral k in let v = frob p u h in shift (a - length v) v | k <- [0 .. a - 1]]

-- derivPolyZ
-- | derivPolyZ : derivative of polynmial P over Fp[x]
derivPolyZ :: Integer -> [Integer] -> [Integer]
derivPolyZ _ [] = []
derivPolyZ p x = init $ reverse [let y = reverse x in 
		let j = fromIntegral i in mods (j * (y!!i)) p | i <- [0..length x - 1]]

-- | squareFreePolyZ
-- squareFreePolyZ p f : gives the euclidean quotient of P and gcd(f,f'). That quotient is a square free polynomial.
squareFreePolyZ :: Integer -> [Integer] -> [Integer]
squareFreePolyZ p x = let y = derivPolyZ p x in 
			let z = trim' $ gcdPolyZ p x y in if z == [1] then x else head $ euclideanPolyMod p x z


-- |berlekamp
-- berlekamp p P: gives a complete factorization of a polynom P of irreducible polynoms over Fp[x].
berlekamp :: Integer -> [Integer] -> [[Integer]]
berlekamp p f = let g = squareFreePolyZ p f in 
	let v = map (trim) $! map (reverse ) $! nullSpaceModIntegerMat p (matrixBerlTranspose p g) in
	let q = length $! v in 
	case q of 1 -> [g]
		  _ -> let (a:b:bs) = phi p g v in 			
			let c = zipgcdPoly p a b in if length c == q then c else zipgcdPoly p c (head bs) where 
zipgcdPoly _ [] _  = []
zipgcdPoly p (x:xs) y = let z = zgcdPoly p x y  in z ++ zipgcdPoly p xs y where
zgcdPoly _ _ [] = []
zgcdPoly p x (y:ys) = let u = gcdPolyZ p x y in if length u > 1 then u:zgcdPoly p x ys else zgcdPoly p x ys where
phi _ _ [] = []
phi p x (v:vs) = let u = filter (\x -> length x > 1) $! [gcdPolyZ p x ((+:) (-) v [i]) | i <- [0 .. p - 1] ] in u : phi p x vs

-- | irreducibilityTestPolyZ
-- irreducibilityTestPolyZ : irreducibility test of polynomials over Fp[x]
irreducibilityTestPolyZ :: Integer -> [Integer] -> Bool
irreducibilityTestPolyZ p f = let g = squareFreePolyZ p f in if g /= f then False else let u = matrixBerlTranspose p g in
	let v = nullSpaceModIntegerMat p u in
	let q = length $ v in 
	case q of 1 -> True
		  _ -> False 
-- | multPoly
-- multPoly : product of polynomials P1, .., Pk in Fp[x].
multPoly p x = head $! multPoly' p x where
multPoly' :: Integer ->[[Integer]]-> [[Integer]] 
multPoly' _ [x] = [x]
multPoly' p (x:t:xs) = let y = multPolyZ p x t in multPoly' p (y:xs)