module Main where

import Stochastic.Distributions
import Stochastic.Analysis
import Stochastic.Uniform
import System.Random
import Stochastic.Tools
import qualified Stochastic.Distributions.Continuous as C
import qualified Stochastic.Distributions.Discrete as D
import Test.HUnit
import Data.Monoid
import Control.Monad
--import Utils
import Control.Monad (unless)

tests = TestList [
  TestLabel "Binomial pmf" bpmf1
  ,TestLabel "ZipF pmf" zipFpmf1
  ,TestLabel "Normal cdf 0" normalCDF0
  ,TestLabel "Normal cdf 1" normalCDF1
  ,TestLabel "Normal cdf 2" normalCDF2
  ,TestLabel "Normal cdf 3" normalCDF3
  ,TestLabel "Normal ChiSquared" chinorm
  ,TestLabel "Exponential ChiSquared" chiexp
  ,TestLabel "Uniform ChiSquared" chiuni
  ,TestLabel "Poisson ChiSquared" chiPoisson
{-  ,TestLabel "Test of grouping" groupTest
--  ,TestLabel "Test histogram" binTest
-}
  ]


--
-- Discrete Tests
--
bpmf1 = TestCase (assertWithinDelta "binomial's pmf must sum to 1" (1e-15) e a)
  where
    b = D.mkBinomial (stdBase 42) 0.4 10
    e = 1
    a = sum [ D.pmf b k | k <- [0..10] ]

zipFpmf1 = TestCase (assertWithinDelta "ZipF's pmf must sum to 1" (1e-15) e a)
  where
    b = D.mkZipF (stdBase 42) 10 1
    e = 1
    a = sum [ D.pmf b k | k <- [1..10] ]

--
-- Continuous Tests
--

stdNormal = C.mkNormal (stdBase 42) 0 1

normalCDF0 = TestCase (assertWithinDelta
                       "Normal CDF at 0 standard deviations"
                       (1e-15) 0.5    $ C.cdf stdNormal 0)
normalCDF1 = TestCase (assertWithinDelta
                       "Normal CDF at 1 standard deviations"
                       (1e-4)  0.8413 $ C.cdf stdNormal 1)
normalCDF2 = TestCase (assertWithinDelta
                       "Normal CDF at 2 standard deviations"
                       (1e-4)  0.9772 $ C.cdf stdNormal 2)
normalCDF3 = TestCase (assertWithinDelta
                       "Normal CDF at 0 standard deviations"
                       (1e-4)  0.9987 $ C.cdf stdNormal 3)

--
-- Stochastic Tests
--

chiTestRandom :: C.Dist
                 -> Double
chiTestRandom g = chiSquaredTest g (mkEmpirical samples) bounds
  where
    (samples, _) = C.rands 1000 g
    hist = fIHistogram samples
    bounds = fmap (lower_bound) hist

chinorm = TestCase (
  assertWithinDelta
  "Sample of the normal distribution passes the ChiSquaredTest with HIGH confidence"
  (2e-1) 0 (chiTestRandom stdNormal))

chiuni = TestCase (
  assertWithinDelta
  "Sample of the uniform distribution passes the ChiSquaredTest with HIGH confidence"
  (1e-1) 0 (chiTestRandom $ C.mkUniform $ stdBase 42))

chiexp = TestCase (
  assertWithinDelta
  "Sample of the exponential distribution passes the ChiSquaredTest with HIGH confidence"
  (0.25) 0 (chiTestRandom $ C.mkExp (stdBase 42) 1))

groupTest = TestCase
            (assertEqual
             "groups where non optimal"
             ([(1,2), (2,3), (3,1)])
             (group [2,1,1,2,2,3])
            )

chiPoisson = TestCase (
  assertWithinDelta
  "Sample of the poisson distribution passes the ChiSquaredTest with HIGH confidence"
  (1.2) 0 (chiTestDiscreteRandom $ D.mkPoisson (stdBase 42) 1))

chiTestDiscreteRandom :: D.Dist
                      -> Double
chiTestDiscreteRandom g = discreteChiSquaredTest g (mkEmpirical samples) bounds
  where
    samples = fmap (toDbl) $ fst $ D.rands 1000 g
    hist = fIHistogram samples
    bounds = fmap (fromDbl . lower_bound) hist
    toDbl = fromInteger . toInteger
    fromDbl = fromInteger . truncate

            
{-
binTest = TestCase
          (assertEqual
           "bins had holes"
           ([])
           (fmap (\x -> (lower_bound x, upper_bound x, frequency x))
            $ fIHistogram samples)
          )
  where
    (samples, _) = C.rands 1000 stdNormal
-}

{-
 let norm = mkNormal (stdBase 42) 0 1 in                                let (samples,_) = rands 1000 norm in                                   let emp = mkEmpirical samples in                                       let hist = fIHistogram samples in                                      let bounds = fmap (lower_bound) hist in                                let chi = chiSquaredTest norm emp bounds in                            let imp = fmap (\x -> (lower_bound x, frequency x, C.cdf norm (lower_bound x), C.cdf emp (lower_bound x)) ) hist in mapM (putStrLn.show) imp-}

pleasantAPIExample :: [(Double, Int, Double)]
pleasantAPIExample =
  let g = stdBase 42 in
  let ([x, y, z], g') = nWayAllocate 20 3 g in
  let g1 = C.mkExp x 1 in
  let g2 = D.mkPoisson y 1 in
  let g3 = C.mkNormal z 0 1 in
  zip3 (randoms g1 :: [Double]) (randoms g2 :: [Int])  (randoms g3 :: [Double]) 
      
      
      


--
-- Custom Assertions
--

assertWithinDelta ::   String  -- ^ The message prefix
               -> Double  -- ^ The maximum difference between expected and actual
               -> Double  -- ^ The expected value
               -> Double  -- ^ The actual value
               -> Assertion
assertWithinDelta preface delta expected actual = 
  unless passed (assertFailure msg)
  where
    passed = (abs (expected - actual) < delta)
    msg = (if null preface then "" else preface ++ "\n") ++
          "expected: " ++ show expected ++ "\n but got: " ++ show actual
  


main :: IO ()
main =
  do
    counts <- runTestTT tests
    return ()