module Stochastic.Analysis(chiSquaredTest, discreteChiSquaredTest) where

import qualified Stochastic.Distributions.Continuous as C
import qualified Stochastic.Distributions.Discrete as D

import Data.List(sort)
import Stochastic.Distributions
import Stochastic.Tools

chiCritD :: (C.ContinuousDistribution g)
            => g
            -> Empirical
            -> Double
            -> Double
chiCritD theory emp =
  chiCrit (degreesOfFreedom emp) (C.degreesOfFreedom theory)

-- k = number of intervals in the empirical histogram
-- s = number of parameters to the theoretical distribution
chiCrit :: Int -> Int -> Double -> Double
chiCrit k s alpha = C.cdf' chi (1-alpha)
  where
    df = k - s - 1
    chi = C.ChiSquared df (stdBase 42)

-- empirical distributions should be
-- given as the second argument
chiSquaredTest :: (C.ContinuousDistribution g)
                  => g
                  -> Empirical
                  -> [Double]
                  -> Double
chiSquaredTest c d sampleAt = if (isNaN final) then (error $ "frog") else final
  where
    final = sum $ fmap f sampleAt
    f x =  let o = cdf d x in
           let e = C.cdf c x in
           if (e == 0)
           then ((e-o)**2)
           else ((e-o)**2) / e


discreteChiSquaredTest :: (D.DiscreteDistribution g)
                          => g
                          -> Empirical
                          -> [Int]
                          -> Double
discreteChiSquaredTest c d sampleAt = sum $ fmap f sampleAt
  where
    f :: Int -> Double
    f 0 = 0
    f x =  let e = D.cdf c x in
           let o = cdf d $ toDbl x in
           if (e == 0)
           then ((e-o)**2)
           else ((e-o)**2) / e
    toDbl :: Int -> Double 
    toDbl = fromInteger . toInteger