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 $ C.cdf c x in
           let e = x in
           let ret = ((e-o)**2) / e in
           if (isNaN ret)
           then error $ (show e) ++ " " ++ (show o) ++ " " ++ (show ret) ++ " " ++ (show (0/0))
           else ret


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