    -- |
    -- Module:      Math.NumberTheory.Recurrences.Linear
    -- Copyright:   (c) 2011 Daniel Fischer
    -- Licence:     MIT
    -- Maintainer:  Daniel Fischer <daniel.is.fischer@googlemail.com>
    --
    -- Efficient calculation of linear recurrent sequences, including Fibonacci and Lucas sequences.

    {-# LANGUAGE BangPatterns     #-}
    {-# LANGUAGE PostfixOperators #-}

    module Math.NumberTheory.Recurrences.Linear
      ( factorial
      , factorialFactors
      , fibonacci
      , fibonacciPair
      , lucas
      , lucasPair
      , generalLucas
      ) where

    import Data.Bits
    import Data.List.Infinite (Infinite(..), (...))
    import qualified Data.List.Infinite as Inf
    import Numeric.Natural
    import Math.NumberTheory.Primes

    -- | Infinite zero-based table of factorials.
    --
    -- >>> take 5 factorial
    -- [1,1,2,6,24]
    --
    -- The time-and-space behaviour of 'factorial' is similar to described in
    -- "Math.NumberTheory.Recurrences.Bilinear#memory".
    factorial :: (Num a, Enum a) => Infinite a
    factorial :: forall a. (Num a, Enum a) => Infinite a
    factorial = (a -> a -> a) -> a -> Infinite a -> Infinite a
    forall b a. (b -> a -> b) -> b -> Infinite a -> Infinite b
    Inf.scanl a -> a -> a
    forall a. Num a => a -> a -> a
    (*) a
    1 (a
    1...)
    {-# SPECIALIZE factorial :: Infinite Int     #-}
    {-# SPECIALIZE factorial :: Infinite Word    #-}
    {-# SPECIALIZE factorial :: Infinite Integer #-}
    {-# SPECIALIZE factorial :: Infinite Natural #-}

    -- | Prime factors of a factorial.
    --
    -- > factorialFactors n == factorise (factorial !! n)
    --
    -- >>> factorialFactors 10
    -- [(Prime 2,8),(Prime 3,4),(Prime 5,2),(Prime 7,1)]
    factorialFactors :: Word -> [(Prime Word, Word)]
    factorialFactors :: Word -> [(Prime Word, Word)]
    factorialFactors Word
    n
      | Word
    n Word -> Word -> Bool
    forall a. Ord a => a -> a -> Bool
    < Word
    2
      = []
      | Bool
    otherwise
      = (Prime Word -> (Prime Word, Word))
    -> [Prime Word] -> [(Prime Word, Word)]
    forall a b. (a -> b) -> [a] -> [b]
    map (\Prime Word
    p -> (Prime Word
    p, Word -> Word
    mult (Prime Word -> Word
    forall a. Prime a -> a
    unPrime Prime Word
    p))) [Prime Word
    forall a. Bounded a => a
    minBound .. Word -> Prime Word
    forall a.
    (Bits a, Integral a, UniqueFactorisation a) =>
    a -> Prime a
    precPrime Word
    n]
      where
        mult :: Word -> Word
        mult :: Word -> Word
    mult Word
    p = Word -> Word -> Word
    go Word
    np Word
    np
          where
            np :: Word
    np = Word
    n Word -> Word -> Word
    forall a. Integral a => a -> a -> a
    `quot` Word
    p
            go :: Word -> Word -> Word
    go !Word
    acc !Word
    x
              | Word
    x Word -> Word -> Bool
    forall a. Ord a => a -> a -> Bool
    >= Word
    p = let xp :: Word
    xp = Word
    x Word -> Word -> Word
    forall a. Integral a => a -> a -> a
    `quot` Word
    p in Word -> Word -> Word
    go (Word
    acc Word -> Word -> Word
    forall a. Num a => a -> a -> a
    + Word
    xp) Word
    xp
              | Bool
    otherwise = Word
    acc

    -- | @'fibonacci' k@ calculates the @k@-th Fibonacci number in
    --   /O/(@log (abs k)@) steps. The index may be negative. This
    --   is efficient for calculating single Fibonacci numbers (with
    --   large index), but for computing many Fibonacci numbers in
    --   close proximity, it is better to use the simple addition
    --   formula starting from an appropriate pair of successive
    --   Fibonacci numbers.
    fibonacci :: Num a => Int -> a
    fibonacci :: forall a. Num a => Int -> a
    fibonacci = (a, a) -> a
    forall a b. (a, b) -> a
    fst ((a, a) -> a) -> (Int -> (a, a)) -> Int -> a
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Int -> (a, a)
    forall a. Num a => Int -> (a, a)
    fibonacciPair
    {-# SPECIALIZE fibonacci :: Int -> Int     #-}
    {-# SPECIALIZE fibonacci :: Int -> Word    #-}
    {-# SPECIALIZE fibonacci :: Int -> Integer #-}
    {-# SPECIALIZE fibonacci :: Int -> Natural #-}

    -- | @'fibonacciPair' k@ returns the pair @(F(k), F(k+1))@ of the @k@-th
    --   Fibonacci number and its successor, thus it can be used to calculate
    --   the Fibonacci numbers from some index on without needing to compute
    --   the previous. The pair is efficiently calculated
    --   in /O/(@log (abs k)@) steps. The index may be negative.
    fibonacciPair :: Num a => Int -> (a, a)
    fibonacciPair :: forall a. Num a => Int -> (a, a)
    fibonacciPair Int
    n
      | Int
    n Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    < Int
    0     = let (a
    f,a
    g) = Int -> (a, a)
    forall a. Num a => Int -> (a, a)
    fibonacciPair (-(Int
    nInt -> Int -> Int
    forall a. Num a => a -> a -> a
    +Int
    1)) in if Int -> Int -> Bool
    forall a. Bits a => a -> Int -> Bool
    testBit Int
    n Int
    0 then (a
    g, -a
    f) else (-a
    g, a
    f)
      | Int
    n Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == Int
    0    = (a
    0, a
    1)
      | Bool
    otherwise = Int -> (a, a)
    forall a. Num a => Int -> (a, a)
    look (Word -> Int
    forall b. FiniteBits b => b -> Int
    finiteBitSize (Word
    0 :: Word) Int -> Int -> Int
    forall a. Num a => a -> a -> a
    - Int
    2)
        where
          look :: Int -> (t, t)
    look Int
    k
            | Int -> Int -> Bool
    forall a. Bits a => a -> Int -> Bool
    testBit Int
    n Int
    k = Int -> t -> t -> (t, t)
    forall {t}. Num t => Int -> t -> t -> (t, t)
    go (Int
    kInt -> Int -> Int
    forall a. Num a => a -> a -> a
    -Int
    1) t
    0 t
    1
            | Bool
    otherwise   = Int -> (t, t)
    look (Int
    kInt -> Int -> Int
    forall a. Num a => a -> a -> a
    -Int
    1)
          go :: Int -> t -> t -> (t, t)
    go Int
    k t
    g t
    f
            | Int
    k Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    < Int
    0       = (t
    f, t
    ft -> t -> t
    forall a. Num a => a -> a -> a
    +t
    g)
            | Int -> Int -> Bool
    forall a. Bits a => a -> Int -> Bool
    testBit Int
    n Int
    k = Int -> t -> t -> (t, t)
    go (Int
    kInt -> Int -> Int
    forall a. Num a => a -> a -> a
    -Int
    1) (t
    ft -> t -> t
    forall a. Num a => a -> a -> a
    *(t
    ft -> t -> t
    forall a. Num a => a -> a -> a
    +t -> t
    forall a. Num a => a -> a
    shiftL1 t
    g)) ((t
    ft -> t -> t
    forall a. Num a => a -> a -> a
    +t
    g)t -> t -> t
    forall a. Num a => a -> a -> a
    *t -> t
    forall a. Num a => a -> a
    shiftL1 t
    f t -> t -> t
    forall a. Num a => a -> a -> a
    + t
    gt -> t -> t
    forall a. Num a => a -> a -> a
    *t
    g)
            | Bool
    otherwise   = Int -> t -> t -> (t, t)
    go (Int
    kInt -> Int -> Int
    forall a. Num a => a -> a -> a
    -Int
    1) (t
    ft -> t -> t
    forall a. Num a => a -> a -> a
    *t
    ft -> t -> t
    forall a. Num a => a -> a -> a
    +t
    gt -> t -> t
    forall a. Num a => a -> a -> a
    *t
    g) (t
    ft -> t -> t
    forall a. Num a => a -> a -> a
    *(t
    ft -> t -> t
    forall a. Num a => a -> a -> a
    +t -> t
    forall a. Num a => a -> a
    shiftL1 t
    g))
    {-# SPECIALIZE fibonacciPair :: Int -> (Int, Int)         #-}
    {-# SPECIALIZE fibonacciPair :: Int -> (Word, Word)       #-}
    {-# SPECIALIZE fibonacciPair :: Int -> (Integer, Integer) #-}
    {-# SPECIALIZE fibonacciPair :: Int -> (Natural, Natural) #-}

    -- | @'lucas' k@ computes the @k@-th Lucas number. Very similar
    --   to @'fibonacci'@.
    lucas :: Num a => Int -> a
    lucas :: forall a. Num a => Int -> a
    lucas = (a, a) -> a
    forall a b. (a, b) -> a
    fst ((a, a) -> a) -> (Int -> (a, a)) -> Int -> a
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Int -> (a, a)
    forall a. Num a => Int -> (a, a)
    lucasPair
    {-# SPECIALIZE lucas :: Int -> Int     #-}
    {-# SPECIALIZE lucas :: Int -> Word    #-}
    {-# SPECIALIZE lucas :: Int -> Integer #-}
    {-# SPECIALIZE lucas :: Int -> Natural #-}

    -- | @'lucasPair' k@ computes the pair @(L(k), L(k+1))@ of the @k@-th
    --   Lucas number and its successor. Very similar to @'fibonacciPair'@.
    lucasPair :: Num a => Int -> (a, a)
    lucasPair :: forall a. Num a => Int -> (a, a)
    lucasPair Int
    n
      | Int
    n Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    < Int
    0     = let (a
    f,a
    g) = Int -> (a, a)
    forall a. Num a => Int -> (a, a)
    lucasPair (-(Int
    nInt -> Int -> Int
    forall a. Num a => a -> a -> a
    +Int
    1)) in if Int -> Int -> Bool
    forall a. Bits a => a -> Int -> Bool
    testBit Int
    n Int
    0 then (-a
    g, a
    f) else (a
    g, -a
    f)
      | Int
    n Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == Int
    0    = (a
    2, a
    1)
      | Bool
    otherwise = Int -> (a, a)
    forall a. Num a => Int -> (a, a)
    look (Word -> Int
    forall b. FiniteBits b => b -> Int
    finiteBitSize (Word
    0 :: Word) Int -> Int -> Int
    forall a. Num a => a -> a -> a
    - Int
    2)
        where
          look :: Int -> (t, t)
    look Int
    k
            | Int -> Int -> Bool
    forall a. Bits a => a -> Int -> Bool
    testBit Int
    n Int
    k = Int -> t -> t -> (t, t)
    forall {t}. Num t => Int -> t -> t -> (t, t)
    go (Int
    kInt -> Int -> Int
    forall a. Num a => a -> a -> a
    -Int
    1) t
    0 t
    1
            | Bool
    otherwise   = Int -> (t, t)
    look (Int
    kInt -> Int -> Int
    forall a. Num a => a -> a -> a
    -Int
    1)
          go :: Int -> t -> t -> (t, t)
    go Int
    k t
    g t
    f
            | Int
    k Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    < Int
    0       = (t -> t
    forall a. Num a => a -> a
    shiftL1 t
    g t -> t -> t
    forall a. Num a => a -> a -> a
    + t
    f,t
    gt -> t -> t
    forall a. Num a => a -> a -> a
    +t
    3t -> t -> t
    forall a. Num a => a -> a -> a
    *t
    f)
            | Bool
    otherwise   = Int -> t -> t -> (t, t)
    go (Int
    kInt -> Int -> Int
    forall a. Num a => a -> a -> a
    -Int
    1) t
    g' t
    f'
              where
                (t
    f',t
    g')
                  | Int -> Int -> Bool
    forall a. Bits a => a -> Int -> Bool
    testBit Int
    n Int
    k = (t -> t
    forall a. Num a => a -> a
    shiftL1 (t
    ft -> t -> t
    forall a. Num a => a -> a -> a
    *(t
    ft -> t -> t
    forall a. Num a => a -> a -> a
    +t
    g)) t -> t -> t
    forall a. Num a => a -> a -> a
    + t
    gt -> t -> t
    forall a. Num a => a -> a -> a
    *t
    g,t
    ft -> t -> t
    forall a. Num a => a -> a -> a
    *(t -> t
    forall a. Num a => a -> a
    shiftL1 t
    g t -> t -> t
    forall a. Num a => a -> a -> a
    + t
    f))
                  | Bool
    otherwise   = (t
    ft -> t -> t
    forall a. Num a => a -> a -> a
    *(t -> t
    forall a. Num a => a -> a
    shiftL1 t
    g t -> t -> t
    forall a. Num a => a -> a -> a
    + t
    f),t
    ft -> t -> t
    forall a. Num a => a -> a -> a
    *t
    ft -> t -> t
    forall a. Num a => a -> a -> a
    +t
    gt -> t -> t
    forall a. Num a => a -> a -> a
    *t
    g)
    {-# SPECIALIZE lucasPair :: Int -> (Int, Int)         #-}
    {-# SPECIALIZE lucasPair :: Int -> (Word, Word)       #-}
    {-# SPECIALIZE lucasPair :: Int -> (Integer, Integer) #-}
    {-# SPECIALIZE lucasPair :: Int -> (Natural, Natural) #-}

    -- | @'generalLucas' p q k@ calculates the quadruple @(U(k), U(k+1), V(k), V(k+1))@
    --   where @U(i)@ is the Lucas sequence of the first kind and @V(i)@ the Lucas
    --   sequence of the second kind for the parameters @p@ and @q@, where @p^2-4q /= 0@.
    --   Both sequences satisfy the recurrence relation @A(j+2) = p*A(j+1) - q*A(j)@,
    --   the starting values are @U(0) = 0, U(1) = 1@ and @V(0) = 2, V(1) = p@.
    --   The Fibonacci numbers form the Lucas sequence of the first kind for the
    --   parameters @p = 1, q = -1@ and the Lucas numbers form the Lucas sequence of
    --   the second kind for these parameters.
    --   Here, the index must be non-negative, since the terms of the sequence for
    --   negative indices are in general not integers.
    generalLucas :: Num a => a -> a -> Int -> (a, a, a, a)
    generalLucas :: forall a. Num a => a -> a -> Int -> (a, a, a, a)
    generalLucas a
    p a
    q Int
    k
      | Int
    k Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    < Int
    0       = [Char] -> (a, a, a, a)
    forall a. HasCallStack => [Char] -> a
    error [Char]
    "generalLucas: negative index"
      | Int
    k Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == Int
    0      = (a
    0,a
    1,a
    2,a
    p)
      | Bool
    otherwise   = Int -> (a, a, a, a)
    look (Word -> Int
    forall b. FiniteBits b => b -> Int
    finiteBitSize (Word
    0 :: Word) Int -> Int -> Int
    forall a. Num a => a -> a -> a
    - Int
    2)
        where
          look :: Int -> (a, a, a, a)
    look Int
    i
            | Int -> Int -> Bool
    forall a. Bits a => a -> Int -> Bool
    testBit Int
    k Int
    i   = Int -> a -> a -> a -> a -> (a, a, a, a)
    go (Int
    iInt -> Int -> Int
    forall a. Num a => a -> a -> a
    -Int
    1) a
    1 a
    p a
    p a
    q
            | Bool
    otherwise     = Int -> (a, a, a, a)
    look (Int
    iInt -> Int -> Int
    forall a. Num a => a -> a -> a
    -Int
    1)
          go :: Int -> a -> a -> a -> a -> (a, a, a, a)
    go Int
    i a
    un a
    un1 a
    vn a
    qn
            | Int
    i Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    < Int
    0         = (a
    un, a
    un1, a
    vn, a
    pa -> a -> a
    forall a. Num a => a -> a -> a
    *a
    un1 a -> a -> a
    forall a. Num a => a -> a -> a
    - a -> a
    forall a. Num a => a -> a
    shiftL1 (a
    qa -> a -> a
    forall a. Num a => a -> a -> a
    *a
    un))
            | Int -> Int -> Bool
    forall a. Bits a => a -> Int -> Bool
    testBit Int
    k Int
    i   = Int -> a -> a -> a -> a -> (a, a, a, a)
    go (Int
    iInt -> Int -> Int
    forall a. Num a => a -> a -> a
    -Int
    1) (a
    un1a -> a -> a
    forall a. Num a => a -> a -> a
    *a
    vna -> a -> a
    forall a. Num a => a -> a -> a
    -a
    qn) ((a
    pa -> a -> a
    forall a. Num a => a -> a -> a
    *a
    un1a -> a -> a
    forall a. Num a => a -> a -> a
    -a
    qa -> a -> a
    forall a. Num a => a -> a -> a
    *a
    un)a -> a -> a
    forall a. Num a => a -> a -> a
    *a
    vn a -> a -> a
    forall a. Num a => a -> a -> a
    - a
    pa -> a -> a
    forall a. Num a => a -> a -> a
    *a
    qn) ((a
    pa -> a -> a
    forall a. Num a => a -> a -> a
    *a
    un1 a -> a -> a
    forall a. Num a => a -> a -> a
    - (a
    2a -> a -> a
    forall a. Num a => a -> a -> a
    *a
    q)a -> a -> a
    forall a. Num a => a -> a -> a
    *a
    un)a -> a -> a
    forall a. Num a => a -> a -> a
    *a
    vn a -> a -> a
    forall a. Num a => a -> a -> a
    - a
    pa -> a -> a
    forall a. Num a => a -> a -> a
    *a
    qn) (a
    qna -> a -> a
    forall a. Num a => a -> a -> a
    *a
    qna -> a -> a
    forall a. Num a => a -> a -> a
    *a
    q)
            | Bool
    otherwise     = Int -> a -> a -> a -> a -> (a, a, a, a)
    go (Int
    iInt -> Int -> Int
    forall a. Num a => a -> a -> a
    -Int
    1) (a
    una -> a -> a
    forall a. Num a => a -> a -> a
    *a
    vn) (a
    un1a -> a -> a
    forall a. Num a => a -> a -> a
    *a
    vna -> a -> a
    forall a. Num a => a -> a -> a
    -a
    qn) (a
    vna -> a -> a
    forall a. Num a => a -> a -> a
    *a
    vn a -> a -> a
    forall a. Num a => a -> a -> a
    - a
    2a -> a -> a
    forall a. Num a => a -> a -> a
    *a
    qn) (a
    qna -> a -> a
    forall a. Num a => a -> a -> a
    *a
    qn)
    {-# SPECIALIZE generalLucas :: Int     -> Int     -> Int -> (Int, Int, Int, Int)                 #-}
    {-# SPECIALIZE generalLucas :: Word    -> Word    -> Int -> (Word, Word, Word, Word)             #-}
    {-# SPECIALIZE generalLucas :: Integer -> Integer -> Int -> (Integer, Integer, Integer, Integer) #-}
    {-# SPECIALIZE generalLucas :: Natural -> Natural -> Int -> (Natural, Natural, Natural, Natural) #-}

    shiftL1 :: Num a => a -> a
    shiftL1 :: forall a. Num a => a -> a
    shiftL1 a
    n = a
    n a -> a -> a
    forall a. Num a => a -> a -> a
    + a
    n