{- |
This module contains some commonly needed functions that makes a Haskell
interpreter an ideal CS calculator. The file is at <Calculator.hs>.
-}

module Calculator
    (
    -- * Display non-negative integer in bases
    hex, oct, bin
    -- * Exponentiation with mod
    , powmod
    -- * Fibonacci numbers
    {- |
       The fibonacci numbers. f0 = 0, f1 = 1, f(n+2) = f(n+1) + f(n).
       We seem to use 0 as the starting index, but notice it is
       compatible with the usual 1-based definition, since f2=1.
     -}
    , fiblist
    , fib
    -- * Primes
    , primelist
    , isprime
    -- * Extended Euclidean algorithm
    , euclid
    -- * Binomial coefficients, combination (nCr)
    , binomial
    , choose
    -- * Ackermann
    -- | The horror!
    , ack
    -- * Continued fractions of square roots
    , cfsqrt
    ) where

import qualified Data.Map as Map
import Numeric
import Data.Char(intToDigit)

-- | non-negative integer to hexadecimal string
hex :: (Integral n, Show n) => n -> String
hex n = showHex n ""

-- | non-negative integer to octal string
oct :: (Integral n, Show n) => n -> String
oct n = showOct n ""

-- | non-negative integer to binary string
bin :: (Integral n, Show n) => n -> String
bin n = showIntAtBase 2 intToDigit n ""

-- | @powmod m b e@ = b^e modulo m
powmod :: (Integral n, Integral e) => n -> n -> e -> n
powmod m b e | m<1 = error "powmod: illegal modulus"
powmod m b e | e<0 = error "powmod: negative exponent"
powmod m b 0 = 1
powmod m b e = let bm = b `mod` m in go bm bm (e-1) where
  x ⊗ y = x*y `mod` m
  -- go z b e = z*b^e
  go z b 0 = z
  go z b e = sq b e where
    sq b e | r==0 = sq (b ⊗ b) q
           | otherwise = go (z ⊗ b) b (e-1)
      where (q,r) = quotRem e 2

-- * Fibonacci numbers

-- | The infinite list of fibonacci numbers.
fiblist :: [Integer]
fiblist = 0 : 1 : zipWith (+) fiblist (tail fiblist)

{- |
   An efficient recursive algorithm for computing /one/ fibonacci number.

   * f (2k+1) = (f k)^2 + (f (k+1))^2

   * f (2k) = 2 * f k * f (k+1) - (f k)^2
 -}
fib :: (Integral b, Num a) => b -> a
fib n = fst (g n)
  where g 0 = (0,1)
        g n | r == 0 = (2*x*y-x2, x2+y2)
            | otherwise = (x2+y2, 2*x*y+y2)
          where (q,r) = quotRem n 2
                (x,y) = g q
                x2 = x*x
                y2 = y*y

-- * Primes

-- | The infinite list of prime numbers from 2 in increasing order.
primelist :: [Integer]
primelist = sieve [2..]
{- This code is from Melissa E. O'Neill,
   "The Genuine Sieve of Eratosthenes",
   http://www.cs.hmc.edu/~oneill/papers/Sieve-JFP.pdf
   For the lack of a common priority queue, we use the suboptimal
   code, the one that uses a map.
-}
sieve xs = sieve' xs Map.empty
    where
      sieve' []     table = []
      sieve' (x:xs) table =
          case Map.lookup x table of
            Nothing    -> x : sieve' xs (Map.insert (x*x) [x] table)
            Just facts -> sieve' xs (foldl reinsert (Map.delete x table) facts)
          where
            reinsert table prime = Map.insertWith (++) (x+prime) [prime] table

-- | Whether a number is prime. This looks for the number in the prime list.
isprime :: Integer -> Bool
isprime n = head (dropWhile (n >) primelist) == n


-- * Extended Euclidean algorithm

-- | @euclid a b@ performs the extended euclidean algorithm to find (d,m,n)
-- such that a*m+b*n = d = gcd(a,b).
euclid :: Integral a => a -> a -> (a, a, a)
euclid a b = euclid' 1 0 a 0 1 b
  where
  euclid' u v w x y z | z==0 = (w, u, v)
                      | otherwise = seq u' $ seq v' $ euclid' x y z u' v' w'
    where (q,w') = quotRem w z
          u' = u-q*x
          v' = v-q*y

{-
Here is how.  Initial state: pair of equations

  a*1 + b*0 = a
  a*0 + b*1 = b

General state and invariant: pair of equations

  a*u + b*v = w            (1)
  a*x + b*y = z            (2)
  gcd(a,b) = gcd(w,z)

Final state:

  a*m + b*n = d
  a*x + b*y = 0
  gcd(a,b) = gcd(d,0) = d

  So we are done, and the first equation is the answer.

How to make progress: move from gcd(w,z) to gcd(z, w mod z).
Namely, subtract (2) from (1) enough number of times, then swap:

  Let q = w div z, r = w mod z.  New state:

  a*x + b*y = z
  a*(u-qx) + b*(v-qy) = w-qz = r

Done!
-}


-- * Binomial coefficients, combination (nCr)

-- For large n, nCr is better computed with Pascal's triangle.

-- | Binomial coefficients of (x+1)^n.
binomial :: (Integral b, Num a) => b -> [a]
binomial 0 = [1]
binomial n = let b = binomial (n-1) in zipWith (+) (b++[0]) (0:b)

-- | nCr.
choose :: (Num a) => Int -> Int -> a
choose n r = binomial n !! (if r<n-r then r else n-r)


-- * Ackermann

-- | Please don't use the Ackermann function!
ack :: (Num a, Eq a) => a -> a -> a
ack 0 n = n+1
ack m 0 = ack (m-1) 1
ack m n = ack (m-1) (ack m (n-1))


-- * Continued fractions of square roots

-- | Continued fraction of square root of n.
cfsqrt :: (Integral a) => a -> [a]
cfsqrt n = f sqrtn sqrtn 1 (n - sqrtn*sqrtn)
    where
      sqrtn = floor (sqrt (fromIntegral n))
      f aj _ _ 0 = [aj]
      f aj bj ck cj = let { ai = (sqrtn + bj) `div` cj
                          ; bi = ai * cj - bj
                          }
                      in aj : f ai bi cj (ck + ai*(bj-bi))