module GNFS.FactorBase
  ( Relation(..)
  , FactorBase(..)
  , buildFactorBase
  , relationRow
  ) where

import Data.Bits (setBit)
import Data.List (foldl')
import qualified Data.Map.Strict as M

import GNFS.FiniteField (rootsModP)
import GNFS.PolySelect (NFSPoly(..))
import GNFS.Polynomial
import GNFS.Util (primesFrom, modInv)
import LinearAlgebra.GF2 (BitVec)
import ModularArithmeticUtils (legendre)
import Primes.ErathosteneSieve (erathosteneSieve)

-- | A relation: the pair (a, b) with a - b*m smooth over the rational base and
-- a - b*alpha smooth over the algebraic base, with the full exponent data.
data Relation = Relation
  { relA   :: !Integer
  , relB   :: !Integer
  , relRat :: [(Integer, Int)]              -- prime -> exponent of |a - b m|
  , relAlg :: [((Integer, Integer), Int)]   -- (p, r) -> valuation of a - b alpha
  } deriving Show

data FactorBase = FactorBase
  { fbPoly      :: NFSPoly
  , fbRat       :: [Integer]                 -- rational base: primes up to the bound
  , fbAlg       :: [(Integer, Integer)]      -- algebraic base: (p, r), f(r) = 0 mod p, simple root
  , fbAlgPrimes :: [Integer]                 -- distinct primes occurring in fbAlg
  , fbAlgLifts  :: [(Integer, [(Integer, Integer)])]
                                             -- per prime p: for each root r, the Hensel lifts
                                             -- (p^k, r_k) for k = 1, 2, ... (lazy, infinite)
  , fbReal      :: [Double]                  -- real roots of f (sign columns)
  , fbQC        :: [(Integer, Integer)]      -- quadratic-character primes (q, s)
  , fbColRat    :: M.Map Integer Int
  , fbColAlg    :: M.Map (Integer, Integer) Int
  , fbNumCols   :: Int
  }

-- | Build the three factor bases (rational, algebraic, quadratic characters).
buildFactorBase :: NFSPoly -> Integer -> Integer -> Int -> FactorBase
buildFactorBase np ratBound algBound nQC = fb
  where
    f    = npCoeffs np
    f'   = derivP f
    rat  = erathosteneSieve ratBound
    -- primes whose roots are all simple; primes with a repeated root are dropped
    algPrimeRoots =
      [ (p, rs)
      | p <- erathosteneSieve algBound
      , let rs = rootsModP p f
      , not (null rs)
      , all (\r -> evalP f' r `mod` p /= 0) rs
      ]
    alg = [ (p, r) | (p, rs) <- algPrimeRoots, r <- rs ]
    qc  = take nQC
            [ (q, s)
            | q <- primesFrom (max algBound ratBound + 1)
            , let rs = rootsModP q f
            , s <- rs
            , evalP f' s `mod` q /= 0
            ]
    reals = realRoots f
    -- Hensel lifting of a simple root r of f mod p to roots mod p^k
    lifts p r0 = go p r0
      where
        go q rk =
          let q'  = q * p
              inv = case modInv (evalP f' rk) q' of
                      Just i  -> i
                      Nothing -> error "Hensel lift: derivative not invertible"
              rk' = (rk - evalP f rk * inv) `mod` q'
          in (q, rk) : go q' rk'
    nRat  = length rat
    nAlg  = length alg
    fb = FactorBase
      { fbPoly      = np
      , fbRat       = rat
      , fbAlg       = alg
      , fbAlgPrimes = map fst algPrimeRoots
      , fbAlgLifts  = [ (p, lifts p r) | (p, r) <- alg ]
      , fbReal      = reals
      , fbQC        = qc
      , fbColRat    = M.fromList (zip rat [0 ..])
      , fbColAlg    = M.fromList (zip alg [0 ..])
      , fbNumCols   = 1 + nRat + nAlg + length reals + length qc
      }

-- | Parity vector of a relation. Column layout:
--   [0] sign of a - b m | rational primes | algebraic ideals | real-place signs | quadratic characters
relationRow :: FactorBase -> Relation -> BitVec
relationRow fb rel = foldl' setBit 0 cols
  where
    a = relA rel
    b = relB rel
    m = npM (fbPoly fb)
    nRat = length (fbRat fb)
    nAlg = length (fbAlg fb)
    offAlg  = 1 + nRat
    offReal = offAlg + nAlg
    offQC   = offReal + length (fbReal fb)
    signCol = [ 0 | a - b * m < 0 ]
    ratCols = [ 1 + fbColRat fb M.! p | (p, e) <- relRat rel, odd e ]
    algCols = [ offAlg + fbColAlg fb M.! pr | (pr, e) <- relAlg rel, odd e ]
    realCols = [ offReal + j | (j, th) <- zip [0 ..] (fbReal fb)
                             , fromInteger a < fromInteger b * th ]
    qcCols  = [ offQC + j | (j, (q, s)) <- zip [0 ..] (fbQC fb)
                          , let v = (a - b * s) `mod` q
                          , legendre v q == q - 1 ]
    cols = signCol ++ ratCols ++ algCols ++ realCols ++ qcCols