module ModularArithmeticUtils (modExp, modMul, legendre, factorOutTwos) where

import Data.Bits (testBit, shiftR)

-- Modular exponentiation: compute a^b mod m
modExp :: Integer -> Integer -> Integer -> Integer
modExp b 0 m = 1
modExp b e m = t * modExp ((b * b) `mod` m) (shiftR e 1) m `mod` m
            where t = if testBit e 0 then b `mod` m else 1

-- Compute (a * b) mod m
modMul :: Integer -> Integer -> Integer -> Integer
modMul a b m = (a * b) `mod` m

-- Compute the Legendre symbol (a/p)
-- (a/p) = a^((p-1)/2) mod p
-- (a/p) in {-1, 0, 1}
legendre :: Integer -> Integer -> Integer
legendre a p = modExp a ((p - 1) `div` 2) p

-- Factor p-1 as q * 2^s, where q is odd
factorOutTwos :: Integer -> (Integer, Integer)
factorOutTwos n = go n 0
  where
    go x s
      | even x    = go (x `div` 2) (s + 1)
      | otherwise = (x, s)

-- jacobi symbol (a/n)
-- (a/n) in {-1, 0, 1}
jacobi :: Integer -> Integer -> Integer
jacobi a n
  | n <= 0         = error "jacobi: n must be positive"
  | even n         = error "jacobi: n must be odd"
  | a `mod` n == 0 = 0
  | otherwise      = go (a `mod` n) n 1
  where
    go :: Integer -> Integer -> Integer -> Integer
    go 0 _ _ = 0
    go 1 _ s = s
    go x m s =
      let (xOdd, e) = factorOutTwos x
          s' = if even e then s else s * jacobi2 m
      in
      if xOdd == 1
        then s'
        else
          let s'' = if (xOdd `mod` 4 == 3) && (m `mod` 4 == 3) then -s' else s'
              xNext = m `mod` xOdd
          in go xNext xOdd s''

-- Jacobi(2, m), helper function
jacobi2 :: Integer -> Integer
jacobi2 m =
  case m `mod` 8 of
    1 -> 1
    7 -> 1
    3 -> -1
    5 -> -1
    _ -> error "jacobi2: unexpected (m mod 8) for odd m"