-- |
-- 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 = forall b a. (b -> a -> b) -> b -> Infinite a -> Infinite b
Inf.scanl 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 forall a. Ord a => a -> a -> Bool
< Word
2
  = []
  | Bool
otherwise
  = forall a b. (a -> b) -> [a] -> [b]
map (\Prime Word
p -> (Prime Word
p, Word -> Word
mult (forall a. Prime a -> a
unPrime Prime Word
p))) [forall a. Bounded a => a
minBound .. 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 forall a. Integral a => a -> a -> a
`quot` Word
p
        go :: Word -> Word -> Word
go !Word
acc !Word
x
          | Word
x forall a. Ord a => a -> a -> Bool
>= Word
p = let xp :: Word
xp = Word
x forall a. Integral a => a -> a -> a
`quot` Word
p in Word -> Word -> Word
go (Word
acc 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 = forall a b. (a, b) -> a
fst forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 forall a. Ord a => a -> a -> Bool
< Int
0     = let (a
f,a
g) = forall a. Num a => Int -> (a, a)
fibonacciPair (-(Int
nforall a. Num a => a -> a -> a
+Int
1)) in if forall a. Bits a => a -> Int -> Bool
testBit Int
n Int
0 then (a
g, -a
f) else (-a
g, a
f)
  | Int
n forall a. Eq a => a -> a -> Bool
== Int
0    = (a
0, a
1)
  | Bool
otherwise = forall a. Num a => Int -> (a, a)
look (forall b. FiniteBits b => b -> Int
finiteBitSize (Word
0 :: Word) forall a. Num a => a -> a -> a
- Int
2)
    where
      look :: Int -> (t, t)
look Int
k
        | forall a. Bits a => a -> Int -> Bool
testBit Int
n Int
k = forall {t}. Num t => Int -> t -> t -> (t, t)
go (Int
kforall a. Num a => a -> a -> a
-Int
1) t
0 t
1
        | Bool
otherwise   = Int -> (t, t)
look (Int
kforall a. Num a => a -> a -> a
-Int
1)
      go :: Int -> t -> t -> (t, t)
go Int
k t
g t
f
        | Int
k forall a. Ord a => a -> a -> Bool
< Int
0       = (t
f, t
fforall a. Num a => a -> a -> a
+t
g)
        | forall a. Bits a => a -> Int -> Bool
testBit Int
n Int
k = Int -> t -> t -> (t, t)
go (Int
kforall a. Num a => a -> a -> a
-Int
1) (t
fforall a. Num a => a -> a -> a
*(t
fforall a. Num a => a -> a -> a
+forall a. Num a => a -> a
shiftL1 t
g)) ((t
fforall a. Num a => a -> a -> a
+t
g)forall a. Num a => a -> a -> a
*forall a. Num a => a -> a
shiftL1 t
f forall a. Num a => a -> a -> a
+ t
gforall a. Num a => a -> a -> a
*t
g)
        | Bool
otherwise   = Int -> t -> t -> (t, t)
go (Int
kforall a. Num a => a -> a -> a
-Int
1) (t
fforall a. Num a => a -> a -> a
*t
fforall a. Num a => a -> a -> a
+t
gforall a. Num a => a -> a -> a
*t
g) (t
fforall a. Num a => a -> a -> a
*(t
fforall a. Num a => a -> a -> a
+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 = forall a b. (a, b) -> a
fst forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 forall a. Ord a => a -> a -> Bool
< Int
0     = let (a
f,a
g) = forall a. Num a => Int -> (a, a)
lucasPair (-(Int
nforall a. Num a => a -> a -> a
+Int
1)) in if forall a. Bits a => a -> Int -> Bool
testBit Int
n Int
0 then (-a
g, a
f) else (a
g, -a
f)
  | Int
n forall a. Eq a => a -> a -> Bool
== Int
0    = (a
2, a
1)
  | Bool
otherwise = forall a. Num a => Int -> (a, a)
look (forall b. FiniteBits b => b -> Int
finiteBitSize (Word
0 :: Word) forall a. Num a => a -> a -> a
- Int
2)
    where
      look :: Int -> (t, t)
look Int
k
        | forall a. Bits a => a -> Int -> Bool
testBit Int
n Int
k = forall {t}. Num t => Int -> t -> t -> (t, t)
go (Int
kforall a. Num a => a -> a -> a
-Int
1) t
0 t
1
        | Bool
otherwise   = Int -> (t, t)
look (Int
kforall a. Num a => a -> a -> a
-Int
1)
      go :: Int -> t -> t -> (t, t)
go Int
k t
g t
f
        | Int
k forall a. Ord a => a -> a -> Bool
< Int
0       = (forall a. Num a => a -> a
shiftL1 t
g forall a. Num a => a -> a -> a
+ t
f,t
gforall a. Num a => a -> a -> a
+t
3forall a. Num a => a -> a -> a
*t
f)
        | Bool
otherwise   = Int -> t -> t -> (t, t)
go (Int
kforall a. Num a => a -> a -> a
-Int
1) t
g' t
f'
          where
            (t
f',t
g')
              | forall a. Bits a => a -> Int -> Bool
testBit Int
n Int
k = (forall a. Num a => a -> a
shiftL1 (t
fforall a. Num a => a -> a -> a
*(t
fforall a. Num a => a -> a -> a
+t
g)) forall a. Num a => a -> a -> a
+ t
gforall a. Num a => a -> a -> a
*t
g,t
fforall a. Num a => a -> a -> a
*(forall a. Num a => a -> a
shiftL1 t
g forall a. Num a => a -> a -> a
+ t
f))
              | Bool
otherwise   = (t
fforall a. Num a => a -> a -> a
*(forall a. Num a => a -> a
shiftL1 t
g forall a. Num a => a -> a -> a
+ t
f),t
fforall a. Num a => a -> a -> a
*t
fforall a. Num a => a -> a -> a
+t
gforall 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 forall a. Ord a => a -> a -> Bool
< Int
0       = forall a. HasCallStack => [Char] -> a
error [Char]
"generalLucas: negative index"
  | Int
k 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 (forall b. FiniteBits b => b -> Int
finiteBitSize (Word
0 :: Word) forall a. Num a => a -> a -> a
- Int
2)
    where
      look :: Int -> (a, a, a, a)
look Int
i
        | forall a. Bits a => a -> Int -> Bool
testBit Int
k Int
i   = Int -> a -> a -> a -> a -> (a, a, a, a)
go (Int
iforall 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
iforall 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 forall a. Ord a => a -> a -> Bool
< Int
0         = (a
un, a
un1, a
vn, a
pforall a. Num a => a -> a -> a
*a
un1 forall a. Num a => a -> a -> a
- forall a. Num a => a -> a
shiftL1 (a
qforall a. Num a => a -> a -> a
*a
un))
        | forall a. Bits a => a -> Int -> Bool
testBit Int
k Int
i   = Int -> a -> a -> a -> a -> (a, a, a, a)
go (Int
iforall a. Num a => a -> a -> a
-Int
1) (a
un1forall a. Num a => a -> a -> a
*a
vnforall a. Num a => a -> a -> a
-a
qn) ((a
pforall a. Num a => a -> a -> a
*a
un1forall a. Num a => a -> a -> a
-a
qforall a. Num a => a -> a -> a
*a
un)forall a. Num a => a -> a -> a
*a
vn forall a. Num a => a -> a -> a
- a
pforall a. Num a => a -> a -> a
*a
qn) ((a
pforall a. Num a => a -> a -> a
*a
un1 forall a. Num a => a -> a -> a
- (a
2forall a. Num a => a -> a -> a
*a
q)forall a. Num a => a -> a -> a
*a
un)forall a. Num a => a -> a -> a
*a
vn forall a. Num a => a -> a -> a
- a
pforall a. Num a => a -> a -> a
*a
qn) (a
qnforall a. Num a => a -> a -> a
*a
qnforall a. Num a => a -> a -> a
*a
q)
        | Bool
otherwise     = Int -> a -> a -> a -> a -> (a, a, a, a)
go (Int
iforall a. Num a => a -> a -> a
-Int
1) (a
unforall a. Num a => a -> a -> a
*a
vn) (a
un1forall a. Num a => a -> a -> a
*a
vnforall a. Num a => a -> a -> a
-a
qn) (a
vnforall a. Num a => a -> a -> a
*a
vn forall a. Num a => a -> a -> a
- a
2forall a. Num a => a -> a -> a
*a
qn) (a
qnforall 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 forall a. Num a => a -> a -> a
+ a
n