-- |
-- Module:      Math.NumberTheory.Primes
-- Copyright:   (c) 2016-2018 Andrew.Lelechenko
-- Licence:     MIT
-- Maintainer:  Andrew Lelechenko <andrew.lelechenko@gmail.com>
--

{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE LambdaCase        #-}
{-# LANGUAGE PostfixOperators  #-}

{-# OPTIONS_GHC -fno-warn-orphans #-}

module Math.NumberTheory.Primes
    ( Prime
    , unPrime
    , toPrimeIntegral
    , nextPrime
    , precPrime
    , UniqueFactorisation(..)
    , factorBack
    , -- * Old interface
      primes
    ) where

import Data.Bits
import Data.Coerce
import Data.List.Infinite (Infinite(..), (...), (....))
import qualified Data.List.Infinite as Inf
import Data.Maybe
import Data.Word
import Numeric.Natural

import Math.NumberTheory.Primes.Counting (nthPrime, primeCount)
import qualified Math.NumberTheory.Primes.Factorisation.Montgomery as F (factorise)
import qualified Math.NumberTheory.Primes.Testing.Probabilistic as T (isPrime)
import Math.NumberTheory.Primes.Sieve.Eratosthenes (primes, sieveRange, primeList, psieveFrom, primeSieve)
import Math.NumberTheory.Primes.Small
import Math.NumberTheory.Primes.Types
import Math.NumberTheory.Utils (toWheel30, fromWheel30)
import Math.NumberTheory.Utils.FromIntegral

-- | A class for unique factorisation domains.
class Num a => UniqueFactorisation a where
  -- | Factorise a number into a product of prime powers.
  -- Factorisation of 0 is an undefined behaviour. Otherwise
  -- following invariants hold:
  --
  -- > abs n == abs (product (map (\(p, k) -> unPrime p ^ k) (factorise n)))
  -- > all ((> 0) . snd) (factorise n)
  --
  -- >>> factorise (1 :: Integer)
  -- []
  -- >>> factorise (-1 :: Integer)
  -- []
  -- >>> factorise (6 :: Integer)
  -- [(Prime 2,1),(Prime 3,1)]
  -- >>> factorise (-108 :: Integer)
  -- [(Prime 2,2),(Prime 3,3)]
  --
  -- This function is a replacement
  -- for 'Math.NumberTheory.Primes.Factorisation.factorise'.
  -- If you were looking for the latter, please import
  -- "Math.NumberTheory.Primes.Factorisation" instead of this module.
  --
  -- __Warning:__ there are no guarantees of any particular
  -- order of prime factors, do not expect them to be ascending. E. g.,
  --
  -- >>> factorise 10251562501
  -- [(Prime 101701,1),(Prime 100801,1)]
  factorise :: a -> [(Prime a, Word)]
  -- | Check whether an argument is prime.
  -- If it is then return an associated prime.
  --
  -- >>> isPrime (3 :: Integer)
  -- Just (Prime 3)
  -- >>> isPrime (4 :: Integer)
  -- Nothing
  -- >>> isPrime (-5 :: Integer)
  -- Just (Prime 5)
  --
  -- This function is a replacement
  -- for 'Math.NumberTheory.Primes.Testing.isPrime'.
  -- If you were looking for the latter, please import
  -- "Math.NumberTheory.Primes.Testing" instead of this module.
  isPrime   :: a -> Maybe (Prime a)

instance UniqueFactorisation Int where
  factorise = coerce . F.factorise
  isPrime n = if T.isPrime (toInteger n) then Just (Prime $ abs n) else Nothing

instance UniqueFactorisation Word where
  factorise = coerce . F.factorise
  isPrime n = if T.isPrime (toInteger n) then Just (Prime n) else Nothing

instance UniqueFactorisation Integer where
  factorise = coerce . F.factorise
  isPrime n = if T.isPrime n then Just (Prime $ abs n) else Nothing

instance UniqueFactorisation Natural where
  factorise = coerce . F.factorise
  isPrime n = if T.isPrime (toInteger n) then Just (Prime n) else Nothing

-- | Restore a number from its factorisation.
factorBack :: Num a => [(Prime a, Word)] -> a
factorBack = product . map (\(p, k) -> unPrime p ^ k)

-- | Smallest prime, greater or equal to argument.
--
-- > nextPrime (-100) ==    2
-- > nextPrime  1000  == 1009
-- > nextPrime  1009  == 1009
nextPrime :: (Bits a, Integral a, UniqueFactorisation a) => a -> Prime a
nextPrime n
  | n <= 2    = Prime 2
  | n <= 3    = Prime 3
  | n <= 5    = Prime 5
  | otherwise = Inf.head $ mapMaybeInf isPrime $
                  Inf.dropWhile (< n) $ fmap fromWheel30 (toWheel30 n ...)
                  -- dropWhile is important, because fromWheel30 (toWheel30 n) may appear to be < n.
                  -- E. g., fromWheel30 (toWheel30 94) == 97

-- | Largest prime, less or equal to argument. Undefined, when argument < 2.
--
-- > precPrime 100 == 97
-- > precPrime  97 == 97
precPrime :: (Bits a, Integral a, UniqueFactorisation a) => a -> Prime a
precPrime n
  | n < 2     = error "precPrime: tried to take `precPrime` of an argument less than 2"
  | n < 3     = Prime 2
  | n < 5     = Prime 3
  | n < 7     = Prime 5
  | otherwise = Inf.head $ mapMaybeInf isPrime $
                  Inf.dropWhile (> n) $ fmap fromWheel30 ((toWheel30 n, toWheel30 n - 1) ....)
                  -- dropWhile is important, because fromWheel30 (toWheel30 n) may appear to be > n.
                  -- E. g., fromWheel30 (toWheel30 100) == 101

mapMaybeInf :: (a -> Maybe b) -> Infinite a -> Infinite b
mapMaybeInf = Inf.foldr . (maybe id (:<) .)

-------------------------------------------------------------------------------
-- Prime sequences

data Algorithm = IsPrime | Sieve

chooseAlgorithm :: Integral a => a -> a -> Algorithm
chooseAlgorithm from to
  | to <= fromIntegral sieveRange
  && to < from + truncate (sqrt (fromIntegral from) :: Double)
  = IsPrime
  | to > fromIntegral sieveRange
  && to < from + truncate (0.036 * sqrt (fromIntegral from) + 40000 :: Double)
  = IsPrime
  | otherwise
  = Sieve

succGeneric :: (Bits a, Integral a, UniqueFactorisation a) => Prime a -> Prime a
succGeneric = \case
  Prime 2 -> Prime 3
  Prime 3 -> Prime 5
  Prime 5 -> Prime 7
  Prime p -> Inf.head $ mapMaybeInf (isPrime . fromWheel30) ((toWheel30 p + 1) ...)

succGenericBounded
  :: (Bits a, Integral a, UniqueFactorisation a, Bounded a)
  => Prime a
  -> Prime a
succGenericBounded = \case
  Prime 2 -> Prime 3
  Prime 3 -> Prime 5
  Prime 5 -> Prime 7
  Prime p -> case mapMaybe (isPrime . fromWheel30) [toWheel30 p + 1 .. toWheel30 maxBound] of
    []    -> error "Enum.succ{Prime}: tried to take `succ' near `maxBound'"
    q : _ -> q

predGeneric :: (Bits a, Integral a, UniqueFactorisation a) => Prime a -> Prime a
predGeneric = \case
  Prime 2 -> error "Enum.pred{Prime}: tried to take `pred' of 2"
  Prime 3 -> Prime 2
  Prime 5 -> Prime 3
  Prime 7 -> Prime 5
  Prime p -> Inf.head $ mapMaybeInf (isPrime . fromWheel30) ((toWheel30 p - 1, toWheel30 p - 2) ....)

-- 'dropWhile' is important, because 'psieveFrom' can actually contain primes less than p.
enumFromGeneric :: Integral a => Prime a -> [Prime a]
enumFromGeneric p@(Prime p')
  = coerce
  $ dropWhile (< p)
  $ concat
  $ takeWhile (not . null)
  $ map primeList
  $ psieveFrom
  $ toInteger p'

smallPrimesLimit :: Integral a => a
smallPrimesLimit = fromIntegral (maxBound :: Word16)

enumFromToGeneric :: (Bits a, Integral a, UniqueFactorisation a) => Prime a -> Prime a -> [Prime a]
enumFromToGeneric p@(Prime p') q@(Prime q')
  | p' <= smallPrimesLimit, q' <= smallPrimesLimit
  = map (Prime . fromIntegral) $ smallPrimesFromTo (fromIntegral p') (fromIntegral q')
  | p' <= smallPrimesLimit
  = map (Prime . fromIntegral) (smallPrimesFromTo (fromIntegral p') smallPrimesLimit)
  ++ enumFromToGeneric' (nextPrime smallPrimesLimit) q
  | otherwise
  = enumFromToGeneric' p q

enumFromToGeneric'
  :: (Bits a, Integral a, UniqueFactorisation a)
  => Prime a
  -> Prime a
  -> [Prime a]
enumFromToGeneric' p@(Prime p') q@(Prime q') = takeWhile (<= q) $ dropWhile (< p) $
  case chooseAlgorithm p' q' of
    IsPrime -> Prime 2 : Prime 3 : Prime 5 : mapMaybe (isPrime . fromWheel30) [toWheel30 p' .. toWheel30 q']
    Sieve   ->
      if q' < fromIntegral sieveRange
      then           primeList $ primeSieve $ toInteger q'
      else concatMap primeList $ psieveFrom $ toInteger p'

enumFromThenGeneric :: (Bits a, Integral a, UniqueFactorisation a) => Prime a -> Prime a -> [Prime a]
enumFromThenGeneric p@(Prime p') (Prime q') = case p' `compare` q' of
  LT -> filter (\(Prime r') -> (r' - p') `rem` delta == 0) $ enumFromGeneric p
    where
      delta = q' - p'
  EQ -> repeat p
  GT -> filter (\(Prime r') -> (p' - r') `rem` delta == 0) $ reverse $ enumFromToGeneric (Prime 2) p
    where
      delta = p' - q'

enumFromThenToGeneric :: (Bits a, Integral a, UniqueFactorisation a) => Prime a -> Prime a -> Prime a -> [Prime a]
enumFromThenToGeneric p@(Prime p') (Prime q') r@(Prime r') = case p' `compare` q' of
  LT -> filter (\(Prime t') -> (t' - p') `rem` delta == 0) $ enumFromToGeneric p r
    where
      delta = q' - p'
  EQ -> if p' <= r' then repeat p else []
  GT -> filter (\(Prime t') -> (p' - t') `rem` delta == 0) $ reverse $ enumFromToGeneric r p
    where
      delta = p' - q'

instance Enum (Prime Integer) where
  toEnum = nthPrime
  fromEnum = integerToInt . primeCount . unPrime
  succ = succGeneric
  pred = predGeneric
  enumFrom = enumFromGeneric
  enumFromTo = enumFromToGeneric
  enumFromThen = enumFromThenGeneric
  enumFromThenTo = enumFromThenToGeneric

instance Enum (Prime Natural) where
  toEnum = Prime . integerToNatural . unPrime . nthPrime
  fromEnum = integerToInt . primeCount . naturalToInteger . unPrime
  succ = succGeneric
  pred = predGeneric
  enumFrom = enumFromGeneric
  enumFromTo = enumFromToGeneric
  enumFromThen = enumFromThenGeneric
  enumFromThenTo = enumFromThenToGeneric

instance Enum (Prime Int) where
  toEnum n = if p > intToInteger maxBound
    then error $ "Enum.toEnum{Prime}: " ++ show n ++ "th prime = " ++ show p ++ " is out of bounds of Int"
    else Prime (integerToInt p)
    where
      Prime p = nthPrime n
  fromEnum = integerToInt . primeCount . intToInteger . unPrime
  succ = succGenericBounded
  pred = predGeneric
  enumFrom = enumFromGeneric
  enumFromTo = enumFromToGeneric
  enumFromThen = enumFromThenGeneric
  enumFromThenTo = enumFromThenToGeneric

instance Bounded (Prime Int) where
  minBound = Prime 2
  maxBound = precPrime maxBound

instance Enum (Prime Word) where
  toEnum n = if p > wordToInteger maxBound
    then error $ "Enum.toEnum{Prime}: " ++ show n ++ "th prime = " ++ show p ++ " is out of bounds of Word"
    else Prime (integerToWord p)
    where
      Prime p = nthPrime n
  fromEnum = integerToInt . primeCount . wordToInteger . unPrime
  succ = succGenericBounded
  pred = predGeneric
  enumFrom = enumFromGeneric
  enumFromTo = enumFromToGeneric
  enumFromThen = enumFromThenGeneric
  enumFromThenTo = enumFromThenToGeneric

instance Bounded (Prime Word) where
  minBound = Prime 2
  maxBound = precPrime maxBound