module TonelliShanks (tonelliShanks) where

import ModularArithmeticUtils (modExp, modMul, legendre)

-- 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)

-- Tonelli-Shanks algorithm: find x such that x^2 = n (mod p)
-- Returns Just x if it exists, Nothing otherwise.
tonelliShanks :: Integer -> Integer -> Maybe Integer
tonelliShanks n p
  | legendre n p /= 1 = Nothing   -- no square root
  | p == 2            = Just n    -- special case
  | otherwise         = Just x
  where
    (q, s) = factorOutTwos (p - 1)

    -- find z which is a quadratic non-residue mod p
    z = head [z' | z' <- [2..p-1], legendre z' p == p - 1]

    m0 = s
    c0 = modExp z q p
    t0 = modExp n q p
    r0 = modExp n ((q + 1) `div` 2) p

    (x, _, _) = loop m0 c0 t0 r0

    loop :: Integer -> Integer -> Integer -> Integer -> (Integer, Integer, Integer)
    loop m c t r
      | t == 0 = (0, c, t)
      | t == 1 = (r, c, t)
      | otherwise =
          let
            i = smallestI 0 t m
            b = modExp c (2 ^ (m - i - 1)) p
            m' = i
            c' = modMul b b p
            t' = modMul t c' p
            r' = modMul r b p
          in loop m' c' t' r'

    -- find smallest i (0 <= i < m) such that t^(2^i) = 1 mod p
    smallestI i t m
      | i >= m = error "no valid i found"
      | modExp t (2 ^ i) p == 1 = i
      | otherwise = smallestI (i + 1) t m