module Maths where

import Data.Ratio
import Text.Printf
import Control.Monad

import qualified Data.List as List
import qualified Data.Char as Char

import qualified Data.Array as Array
import Numeric

-- fast prime sieve imports
import Data.Array.IO
import Data.Array.Base
import Data.Bits

-- |Answer all subsets of a set.
powerset          :: [a] -> [[a]]
powerset []       =  [[]]
powerset (x:xs)   =  (powerset xs) ++ (map (x:) (powerset xs))

-- |TODO: figure out why this works
powersetVoodoo  = filterM (const [True, False])

-- |Answer a list of lists with the first arg interspersed at all
-- positions in the list arg.
-- E.g., everywhere 1 [4,5] yields [[1,4,5],[4,1,5],[4,5,1]].
-- Copied from <http://www.haskell.org/hawiki/LicensedPreludeExts>.
everywhere :: a -> [a] -> [[a]]
everywhere x []     = [[x]]
everywhere x (y:ys) = (x:y:ys) : [ y:zs | zs <- everywhere x ys ]

-- |Answer all permutations of the given list.
-- Copied from <http://www.haskell.org/hawiki/LicensedPreludeExts>.
permutations :: [a] -> [[a]]
permutations []     = [[]]
permutations (x:xs) = [ zs | ys <- permutations xs , zs <- everywhere x ys ]

-- |Answer the number of k-permutations from a collection of n objects.
permutations_count :: Integral a => a -> a -> a
permutations_count n k = factorial n `div` factorial (n-k)

-- |Answer a list of all combinations (of any length) of the given list.
-- Copied from <http://www.haskell.org/hawiki/LicensedPreludeExts>.
combinations :: [a] -> [[a]]
combinations [] = [[]]
combinations (x:xs)
    = combinations xs ++ [ x:xs' | xs' <- combinations xs ]

-- |Answer the number of ways to choose k unordered objects from a collection
-- of n objects.
combinations_count :: Integral a => a -> a -> a
combinations_count n k 
  = if k == 0 || k == n
       then 1
       else factorial n `div` (factorial k * factorial (n - k))

-- |Answer the probability of an event with probability p occurring r times
-- in n trials.
binomial       :: (Integral a) => a -> a -> Ratio a -> Ratio a
binomial n r p =  (num_combs * (p^r) * (((1%1)-p) ^ (n - r)))
  where num_combs = (combinations_count n r) % 1

-- |Calculate the binomial value with `binomial` and print the result
-- as a double with 2 decimal places.
printBinomial      :: Integral a => a -> a -> Ratio a -> IO ()
printBinomial n r p = printf "%0.2f\n" (ratioToFloating (binomial n r p) :: Double)

-- |Calculate the binomial value assuming that the probability of the event
-- occurring is 1\/n.
binomialSimple :: Integral a => a -> a -> Ratio a
binomialSimple n r = binomial n r (1%n)

-- |Calculate the binomial value with 'binomialSimple' and print the
-- result as a Double with 2 decimal places.
printBinomialSimple     :: Integral a => a -> a -> IO ()
printBinomialSimple n r =  printBinomial n r (1%n)

-- |Answer the given rational as a floating approximation.
ratioToFloating   :: (Integral a, Floating b) => Ratio a -> b
ratioToFloating x = (fromIntegral $ numerator x) / (fromIntegral $ denominator x)

-- |Calculate the length of the farey sequence with denominator n, including 
-- the value for 0\/n and n\/n (so subtract 2 if you want the more standard
-- value excluding these).
fareylen :: Integral a => a -> a
fareylen n = 1 + f 1 0
  where f x acc | x == n    = term x acc
                | otherwise = f (x+1) (term x acc)
        term x acc = acc + totient x

-- |Calculate the length of the farey sequence with denominator n, including 
-- the value for 0\/n and n\/n (so subtract 2 if you want the more standard
-- value excluding these). This version is optimized for n small enough that
-- totient n fits in an Int. 
fareylenf :: Integral a => a -> a
fareylenf n = 1 + (fromIntegral $ f 1 0)
  where f x acc | x == n    = term (fromIntegral x) acc
                | otherwise = f ((fromIntegral x)+1) (term (fromIntegral x) acc)
        term :: Int -> Integer -> Integer
        term x acc = acc + fromIntegral (totient x)

-- |Euler's totient function, which calculates the number of positive
-- integers <=n that are relatively prime to (i.e., do not contain any factor 
-- in common with) n, where 1 is counted as being relatively prime to all 
-- numbers. In number theory, the totient phi(n) of a positive integer n 
-- is defined to be the number of positive integers less than or equal to n and
-- coprime to n. 
-- <http://en.wikipedia.org/wiki/Totient>
-- <http://mathworld.wolfram.com/TotientFunction.html>
totient :: Integral a => a -> a
totient 1         = 1
totient n | n > 1 = round $ foldr (*) (fromIntegral n) (t [head p | p <- primefactors])
  where primefactors = List.group $ factors n
        t []         = []
        t (x:xs)     = 1 - (1/(fromIntegral x)) : t xs
        
-- |Creates an infinite list of n-gon numbers; e.g., for n = 3, returns
-- the triangular numbers (1, 3, 6, 10..), and for n = 5, returns the pentagon 
-- numbers (1, 5, 12, 22..).                    
polyx :: Integral a => a -> [a]
polyx n = polyx' 0 1 diff
  where diff = n-2
        polyx' curr incr diff = next : polyx' next (incr + diff) diff
          where next = curr + incr
          
-- |Convert an integral number to its binary equivalent as a string.
integralToBinary :: Integral a => a -> String
integralToBinary n | n < 0      = error "Arg must be non-negative, not " ++ show n
                   | otherwise  = (showIntAtBase 2 (Char.intToDigit) n) ""

-- |Calculate the factorial of a positive integer.
factorial :: Integral a => a -> a
factorial n | n < 0     = error "Integer must be non-negative."
            | n == 0    = 1
            | otherwise = product [1..n]

-- |Array of factorial results of the 10 digits, from 0 to 9.
factorials :: Array.Array Int Int
factorials = Array.listArray (0,9) (1:1:[product [2..n] | n<- [2..9]])

-- |Calculate the nth triangle number (1, 3, 6, 10..).
tri :: Integral a => a -> a
tri n = round (fromIntegral (n * (n+1)) / 2)

-- |A stream of triangle numbers.
tris :: Integral a => [a]
tris = polyx 3

-- |A stream of pentagon numbers.
pents :: Integral a => [a]
pents = polyx 5

-- |An infinite list of fibonnaci numbers
fibs :: Integral a => [a]
fibs = 1 : 1 : [ a + b | (a, b) <- zip fibs (tail fibs)]

-- A non-recursive implementation of Fibonacci calculation that takes
-- advantage of the closed-form for a given Fibonacci number.
-- DOESN'T WORK FOR LARGE N (>100) DUE TO ROUND-OFF ERROR
--fib :: Integral a => a -> a
--fib n = floor $ (((phi^n - psi^n) / root5) :: Double)
--  where root5 = (sqrt 5)          :: Double
--        phi   = ((1 + root5) / 2) :: Double
--        psi   = ((1 - root5) / 2) :: Double

-- |An infinite list of prime numbers, using a simple implementation of the
-- sieve of Eratosthenes.
primes :: Integral a => [a]
primes = 2 : sieve [3, 5..]
  where sieve (p:rest) = p : sieve [n | n <- rest, mod n p /= 0]
  
-- |A finite list of prime numbers from 2 to n (inclusive), using a simple 
-- implementation of the sieve of Eratosthenes.
primesn :: Integral a => a -> [a]
primesn n = 2 : sieve [3,5..n]
  where sieve [] = []
        sieve (p:ps) = p : sieve [x | x <- ps, mod x p /= 0]
             
-- |Determine if d divides into n evenly.
divides    :: Integral a => a -> a -> Bool
divides d n = rem n d == 0

-- |Determine the least natural number greater than 1 that divides n.
ld         :: Integral a => a -> a
ld n        = ldf 2 n

-- |Determine the least natural number greater than or equal to k that that divides n.
ldf        :: Integral a => a -> a -> a
ldf k n     | rem n k == 0 = k
	        | k*k > n      = n
	        | otherwise    = ldf (k+1) n

-- |Determine if the given integral argument is a prime number.
isprime :: Integral a => a -> Bool
isprime n = n > 1 && (n == 2 || n == 3 || ((modn6 == 1 || modn6 == 5) && n == ldf' 3 n))
  where modn6 = mod n 6
        ldf' k n | divides k n = k
                 | k*k > n     = n
                 | otherwise   = ldf' (k+2) n

-- |Determine the unique prime factorization of a positive integer.
factors :: Integral a => a -> [a]
factors 1 = []
factors n = p : factors (div n p) where p = ld n

-- |Determine the unique prime factors of a positive integer.
ufactors :: Integral a => a -> [a]
ufactors = List.nub . factors

-- |Calculate the proper divisors of n (divides n evenly and less than n). 
divisors :: Integral a => a -> [a]
divisors n = [x | x <- [1..n-1], mod n x == 0]

-- |Calculate the number of divisors of n, including 1 and n itself.
-- This implementation relies on the fact that the number of divisors for a 
-- number whose prime factorization is (p1^a) (p2^b) (p3^c)... (px^n) is equal
-- to (a+1) (b+1) (c+1) ... (n+1).
numdivisors :: Integral a => a -> Int
numdivisors n = let ps = List.group $ factors n
                in product [1 + length p | p <- ps]

-- |Calculate the sum of the divisors of a number.
-- For example: σ(72) = 15×13 = 195.
sumdivisors :: Integral a => a -> a
sumdivisors n = let primefactors = List.group $ factors n
                in product [pnumdivisors (head p) (fromIntegral $ length p) | p <- primefactors]

-- |Calculate the number of divisors of the given prime number raised to the given
-- power, using the formula (p^(x+1)-1)\/(p-1).
pnumdivisors :: Integral a => a -> a -> a
pnumdivisors p x = let numer = fromIntegral (p^(x+1) - 1)
                       denom = fromIntegral p-1
                   in round $ numer / denom
       
-- |Determine if 2 integrals are coprime (have gcd of 1).
coprime :: Integral a => a -> a -> Bool   
coprime a b = 1 == gcd a b

-- |Determine the arithmetic mean of a list of one or more numbers. 
-- If the list is empty, 0 is returned.
mean :: Fractional a => [a] -> a
mean [] = 0
mean xs = (sum xs) / (fromIntegral $ length xs)

-- |Approximate the sum of the first n primes, using the formula (Bach and 
-- Shallit, 1996): sigma n ~ 1\/2 * n^2 * ln n
sum_of_n_primes_app :: Integral a => a -> Float
sum_of_n_primes_app n = (1/2) *  fromIntegral (n*n) * log (fromIntegral n)

-- |Determine the sum of the first n primes. This is very slow. Consider the
-- function that approximates the sum, sum_of_n_primes_app.
sum_of_n_primes :: Integral a => a -> Float
sum_of_n_primes n = fromIntegral $ sum $ take (fromIntegral n) primes

-- A very fast implementation of a prime generation algorithm; generates
-- all primes less than the argument.
-- Not sure where I found this function originally. I thought I got
-- it from the language shootout, but there doesn't appear to be a prime
-- number generation task (anymore?).
-- Not included in Haddock docs, because Haddock merges this comment
-- in with modexp if there is no signature declared for primesf and
-- it is unable to parse the signature if we do include it.
primesf m =  do ps <- newArray (2,m) True >>= \a -> loop a m (2::Int)
                return ps
 where 
  loop arr m n | arr `seq` m `seq` n `seq` False = undefined
  loop arr m n = if n == m then return [] else do
      el <- unsafeRead arr n
      if el then let loop' j | j > m     = do a <- loop arr m (n+1); return (n : a) 
                             | otherwise = unsafeWrite arr j False >> loop' (j+n)
                 in loop' (n+n)
            else loop (arr :: IOUArray Int Bool) m (n+1)
            
-- |Calculates x^y modulo N. This algorithm has complexity O(n^3) where
-- n is the size in bits of the largest of x, y, and N.
modexp :: Integer -> Integer -> Integer -> Integer
modexp x y n =
  if y == 0
     then 1
   else
     let z = modexp x (y `div` 2) n
     in case even y of
          True  -> (z*z) `mod` n
          _     -> (x * z * z) `mod` n

-- |A formula for computing e, from Courant\/Fritz;
-- it converges quite fast, and already is perfectly
-- accurate to limits of Double at e_18.
-- e = 1\/0! + 1\/1! + 1\/2! + ... + 1\/n!
e_calc :: Int -> Double
e_calc n = sum $ take n $ map (1/) $ scanl (*) 1 [1..]

-- |Another series for computing e from Courant/Fritz,
-- but this one is very slow to converge: at e_1000
-- it is still only 2.7169.
e_calc' :: Int -> Double
e_calc' n =  (1 + 1/n') ** n'
  where n' = fromIntegral n