{-# LANGUAGE TypeFamilies #-}
module Stochastic.Tools where

import Control.Monad.State.Lazy
import Data.Maybe
import qualified Data.Map as Map
import Data.List(sort)

maybeHead :: [a] -> Maybe a
maybeHead []     = Nothing
maybeHead (x:xs) = Just x

headOrElse :: a -> [a] -> a
headOrElse d ls = fromMaybe d (maybeHead ls)

mapTuple :: (a -> b) -> (c -> d) -> (a, c) -> (b, d)    
mapTuple f g (x, y) = (f x, g y)

histogram :: Double -> [Double] -> [Int]
histogram precision ls = map lookup0 [0..limit]
  where
    limit = truncate $ (1.0 / precision)
    divs = map (\x -> x / precision) ls
    ints = map (truncate) divs
    m   = foldl
          (\b a ->
            Map.insert a ((fromMaybe 0 (Map.lookup a b))+1) b)
          Map.empty
          ints
    lookup0 = (\x -> fromMaybe 0 $ Map.lookup x m)

comb :: Int -> Int -> Double
comb n k =
  (fromInteger $ foldr (*) 1 ls) / (fromInteger $ fac (n-k))
  where
    ls = [(toInteger k)..(toInteger n)]


gamma :: Double -> Double
gamma = stirlingsApprox

lower_incomplete_gamma :: Double -> Double -> Double
lower_incomplete_gamma = estimate_lower_gamma 10

estimate_lower_gamma n s x = (x**s)*(gamma s) * (exp (-x)) * approx
  where
    approx = sum $ fmap term [0..n]
    term k = (x**k) / (gamma (s + k + 1))

stirlingsApprox n 
  | n < 0.5   = stirlingsApprox (1-n)
  | otherwise = (sqrt (2*pi*n)) * (exp $ n * ((log n) - 1))

fac :: Int -> Integer
fac n = head $ drop (n) factorial

factorial :: [Integer]
factorial = 1 : (acc (*) 1 [1..])

acc :: (a -> a -> a) -> a -> [a] -> [a]
acc c z ls = f z ls
  where
    f y [] = []
    f y (x:xs) = (x `c` y) : (f (x `c` y) xs)

fib :: Int -> Integer
fib n = head $ drop (n-1) fibinacci

fibinacci :: [Integer]
fibinacci = 1 : 1 : (zipWith (+) fibinacci (tail fibinacci))

harmonics :: Double -> [Double]
harmonics s = (h 1)
  where
    h n = (1.0/(n**s)) : h (n+1)

statefully :: (g -> (a, g)) -> Int -> g -> ([a], g)
statefully f n g0 = case n of
    0 -> ([], g0)
    x -> (r:rest, g2)
      where
        (r, g1) = f g0
        (rest, g2) = statefully f (n-1) g1

statefullyTakeWhile :: (g -> (a, g)) ->
                       ([a] -> Bool) ->
                       g ->
                       ([a], g)
statefullyTakeWhile f p g0 = r ([], g0)
  where
    r (list, g1)
      | p list    = (list, g1)
      | otherwise = mapTuple (\l -> l:list) (id) (f g1)

type Histogram = [Datagram]
data Datagram = Datagram {
  lower_bound :: Double,
  upper_bound :: Double,
  frequency :: Int,
  rel_frequence :: Double,
  cum_frequence :: Double,
  slope :: Double
  } deriving (Eq)

instance Show Datagram where
  show d = show (lower_bound d, upper_bound d, rel_frequence d)

--peice wise linear 
fIHistogram :: [Double] -> Histogram
fIHistogram list =
  datagramFromRaw len iSize $ bin (lowerOf iSize) list
  where
    toDbl = fromInteger . toInteger
    len = length list
    iSize :: Double
    iSize = (maxf list -
             minf list) /
            iCount
    iCount = sqrt (toDbl $ len)

maxf (x:xs) = foldl (max) x xs
minf (x:xs) = foldl (min) x xs

-- TODO fix 'holes'
datagramFromRaw :: Int -> Double -> [(Double, Int)] -> [Datagram]
datagramFromRaw count iSize = accMap f z
  where
    f acc (x, c)  = Datagram {
      lower_bound = x,
      upper_bound = x + iSize,
      frequency = c,
      rel_frequence = rf,
      cum_frequence = rf + (cum_frequence acc),
      slope = rf / iSize
      }
      where rf = (toDbl c / toDbl count)
    z (x, c) = Datagram {
      lower_bound = x,
      upper_bound = x + iSize,
      frequency = c,
      rel_frequence = rf,
      cum_frequence = rf,
      slope = rf / iSize
      }
      where rf = (toDbl c / toDbl count)
    toDbl = fromInteger . toInteger

accMap :: (b -> a -> b) -> (a -> b) -> [a] -> [b]
accMap f z ls = g (z $ head ls) (tail ls)
  where
    g prev []     = [prev]
    g prev (x:xs) = prev : (g new xs)
      where new = (f prev x)
{-
Takes the set 
w : x : y : z : []
and produces the set 
(f w x) : (f x y) : (f y z) : []
-}    
biFold :: (a -> a -> b) -> [a] -> [b]
biFold f = g
  where
    g []  = []
    g [x] = []
    g (x:y:ys) = (f x y) : (g (y:ys))

dependentMap :: (a -> a -> b) -> (a -> b) -> [a] -> [b]
dependentMap f z xs = (z $ head xs) : (biFold f xs)
    
    
bin :: (Double -> Double) -> [Double] -> [(Double, Int)]
bin lowerOf ls = group $ map (lowerOf) ls

lowerOf :: Double -> Double -> Double
lowerOf iSize v = iSize * (fromInteger $ truncate (v / iSize))


    
group :: (Ord a) => [a] -> [(a, Int)]
group = groupSeq . sort
    
groupSeq :: (Eq a) => [a] -> [(a, Int)]
groupSeq list = foldr g [] list
  where
    g x [] = [(x, 1)]
    g x ((y,c):ys)
      | x == y    = (y, c+1):ys
      | otherwise = (x, 1):(y,c):ys