-- |
-- 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 #-}

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

import Data.Bits
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) => [a]
factorial :: [a]
factorial = (a -> a -> a) -> a -> [a] -> [a]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl a -> a -> a
forall a. Num a => a -> a -> a
(*) a
1 [a
1..]
{-# SPECIALIZE factorial :: [Int]     #-}
{-# SPECIALIZE factorial :: [Word]    #-}
{-# SPECIALIZE factorial :: [Integer] #-}
{-# SPECIALIZE factorial :: [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 :: 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 :: 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 -> (a, a)
look Int
k
        | Int -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Int
n Int
k = Int -> a -> a -> (a, a)
forall a. Num a => Int -> a -> a -> (a, a)
go (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) a
0 a
1
        | Bool
otherwise   = Int -> (a, a)
look (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
      go :: Int -> a -> a -> (a, a)
go Int
k a
g a
f
        | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0       = (a
f, a
fa -> a -> a
forall a. Num a => a -> a -> a
+a
g)
        | Int -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Int
n Int
k = Int -> a -> a -> (a, a)
go (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) (a
fa -> a -> a
forall a. Num a => a -> a -> a
*(a
fa -> a -> a
forall a. Num a => a -> a -> a
+a -> a
forall a. Num a => a -> a
shiftL1 a
g)) ((a
fa -> a -> a
forall a. Num a => a -> a -> a
+a
g)a -> a -> a
forall a. Num a => a -> a -> a
*a -> a
forall a. Num a => a -> a
shiftL1 a
f a -> a -> a
forall a. Num a => a -> a -> a
+ a
ga -> a -> a
forall a. Num a => a -> a -> a
*a
g)
        | Bool
otherwise   = Int -> a -> a -> (a, a)
go (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) (a
fa -> a -> a
forall a. Num a => a -> a -> a
*a
fa -> a -> a
forall a. Num a => a -> a -> a
+a
ga -> a -> a
forall a. Num a => a -> a -> a
*a
g) (a
fa -> a -> a
forall a. Num a => a -> a -> a
*(a
fa -> a -> a
forall a. Num a => a -> a -> a
+a -> a
forall a. Num a => a -> a
shiftL1 a
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 :: 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 :: 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 a. Num a => Int -> a -> a -> (a, a)
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 :: 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 :: a -> a
shiftL1 a
n = a
n a -> a -> a
forall a. Num a => a -> a -> a
+ a
n