module GNFS.PolySelect
  ( NFSPoly(..)
  , selectPolynomial
  ) where

import Data.Maybe (listToMaybe)
import GNFS.FiniteField (irreducibleModP)
import GNFS.Polynomial
import GNFS.Util (iroot)

-- | A monic polynomial f of degree d and an integer m with f(m) = n.
data NFSPoly = NFSPoly
  { npDegree :: Int
  , npCoeffs :: Poly      -- ^ little-endian, monic
  , npM      :: Integer
  } deriving Show

-- | Base-m polynomial selection. Returns a monic polynomial of the requested degree, irreducible over Q (certified by irreducibility modulo some small prime), with f(m) = n.
-- Balanced digits are tried first (smaller coefficients => smaller norms).
selectPolynomial :: Int -> Integer -> Maybe NFSPoly
selectPolynomial d n = listToMaybe
  [ NFSPoly d f m
  | m <- takeWhile (\m' -> m' > 1 && 2 * m' ^ d > n) [m0, m0 - 1 ..]
  , digits <- [balanced m (plainDigits m n), plainDigits m n]
  , let f = trimP digits
  , length f == d + 1
  , last f == 1
  , irreducibleSomewhere f
  ]
  where
    m0 = iroot d n
    smallPrimes = [3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73]
    irreducibleSomewhere f = any (\p -> irreducibleModP p f) smallPrimes

plainDigits :: Integer -> Integer -> [Integer]
plainDigits _ 0 = []
plainDigits m x = let (q, r) = x `divMod` m in r : plainDigits m q

-- | Convert digits in [0, m) to balanced digits in (-m/2, m/2], propagating carries.
balanced :: Integer -> [Integer] -> [Integer]
balanced m = go 0
  where
    go 0 []     = []
    go c []     = [c]
    go c (x:xs) =
      let v = x + c
      in if 2 * v > m then (v - m) : go 1 xs else v : go 0 xs