module Stochastic.ZipF (mkZipF, ZipF) where

import Helpers
import Stochastic.Uniform
import Data.Maybe
import Stochastic.Distribution

data ZipF = ZipF Int Double Uniform

mkZipF :: Uniform -> Int -> Double -> ZipF
mkZipF base n slope = ZipF n slope base

-- index, number, cumlative
harmonics :: Double -> [(Double, Double, Double)]
harmonics s = (h 1 0)
  where
    h n acc = (n, v, v+acc) : h (n+1) (v+acc)
      where
        v = (1.0/(n**s))

toDbl = fromInteger . toInteger

f n s d = h2 $ h1 hs
  where
    hs      = harmonics s
    mx      = _3 $ head $ drop (n-1) $ take n $ hs
    h1 xs    = headOrElse (toDbl n, 0, 0) $
              (dropWhile (\(i, v, c) -> c < (d * mx))) (take (n) xs)
    h2 (x, _, _) = truncate x
  
g = f 10 1

instance DiscreteDistribution ZipF where
  randIntIn (a, b) (ZipF n slope u0) = (f n slope d, ZipF n slope u1)
    where
      (d, u1) = randDouble u0

_1 :: (a, b, c) -> a
_1 (x, y, z) = x

_3 :: (a, b, c) -> c
_3 (x, y, z) = z