-----------------------------------------------------------------------------
-- Copyright 2020, Ideas project team. This file is distributed under the
-- terms of the Apache License 2.0. For more information, see the files
-- "LICENSE.txt" and "NOTICE.txt", which are included in the distribution.
-----------------------------------------------------------------------------
module Domain.Statistics.Rules
   ( addNRule
   , addAverageRule
   , addVarianceRule
   , addStandardDeviationRule
   , addStandardErrorRule
   ) where

import Control.Monad
import Domain.Math.Expr
import Domain.Statistics.ComponentSet
import Ideas.Common.Library
import Domain.Statistics.Data

----------------------------------------------------------
-- Rules on the data component

addNRule :: Rule ComponentSet
addNRule = describe "Rule for computing the n of the dataset" $ makeRule "component.n" f
 where
   f :: ComponentSet -> Maybe ComponentSet
   f cs = do guard (not $ contains cs SampleSize)
             sample <- getData DataSet cs
             let nCs = nComponents sample
             return $ foldr (\(nm, c) cs' -> append nm c cs') cs nCs
   nComponents (SingleSample xs)      = [(SampleSize,     CExpr $ Nat $ toInteger (length xs))]
   nComponents (TwoSample    xs ys _) = [(One SampleSize, CExpr $ Nat $ toInteger (length xs)),
                                         (Two SampleSize, CExpr $ Nat $ toInteger (length ys))]

addAverageRule :: Rule ComponentSet
addAverageRule = describe "Rule for adding the formula for computing the mean" $ makeRule "component.mean" f
 where
   f :: ComponentSet -> Maybe ComponentSet
   f cs = do guard (not $ contains cs SampleMean)
             guard (not $ contains cs (One SampleMean))
             guard (not $ contains cs (Two SampleMean))
             sample   <- getData DataSet    cs
             -- n        <- getExpr SampleMean cs
             case sampleMean sample of
                [mean]         -> return $ append SampleMean       (CExpr $ fromDouble mean) cs
                [mean0, mean1] -> return $ append (Two SampleMean) (CExpr $ fromDouble mean1) (append (One SampleMean) (CExpr $ fromDouble mean0) cs)
                _              -> Nothing


addVarianceRule :: Rule ComponentSet
addVarianceRule = describe "Rule for adding the formula for computing the variance" $ makeRule "component.variance" f
 where
   f :: ComponentSet -> Maybe ComponentSet
   f cs = do guard (not $ contains cs SampleVariance)
             guard (not $ contains cs (One SampleVariance))
             guard (not $ contains cs (Two SampleVariance))
             sample  <- getData DataSet    cs
             -- n       <- getExpr SampleSize cs
             -- mean    <- getExpr SampleMean cs
             test    <- getTestType TestChoice cs
             -- s = sum((X_i - avg(X))^2) / n
             case sampleVariance (test /= ZTest) sample of
                [var]        -> return $ append SampleVariance       (CExpr $ fromDouble var) cs
                [var0, var1] -> return $ append (Two SampleVariance) (CExpr $ fromDouble var1) (append (One SampleVariance) (CExpr $ fromDouble var0) cs)
                _            -> Nothing

addStandardDeviationRule :: Rule ComponentSet
addStandardDeviationRule = describe "Rule for adding the formula for computing the deviation" $ makeRule "component.standard-deviation" f
 where
   f :: ComponentSet -> Maybe ComponentSet
   f cs = do guard (not $ contains cs SampleSdev)
             guard (not $ contains cs (One SampleSdev))
             guard (not $ contains cs (Two SampleSdev))
             sample  <- getData DataSet    cs
             -- n       <- getExpr SampleSize cs
             -- mean    <- getExpr SampleMean cs
             test    <- getTestType TestChoice cs
             -- s = sqrt(sum((X_i - avg(X))^2) / n)
             case map sqrt $ sampleVariance (test /= ZTest) sample of
                [sd]       -> return $ append SampleSdev       (CExpr $ fromDouble sd) cs
                [sd0, sd1] -> return $ append (Two SampleSdev) (CExpr $ fromDouble sd1) (append (One SampleSdev) (CExpr $ fromDouble sd0) cs)
                _          -> Nothing

addStandardErrorRule :: Rule ComponentSet
addStandardErrorRule = describe "Rule for adding the formula for computing the standard error" $ makeRule "component.standard-error" f
 where
   f :: ComponentSet -> Maybe ComponentSet
   f cs = do guard (not $ contains cs StandardError)
             guard (not $ contains cs (One StandardError))
             guard (not $ contains cs (Two StandardError))
             sd <- getExpr SampleSdev cs
             -- n  <- getExpr SampleSize  cs
             -- SE = s / sqrt(n)
             return $ append StandardError (CExpr $ sd / sqrt(Var "n")) cs